Skip to content

Commit 34f301f

Browse files
rgommersbetatim
rgommers
and
betatim
authored
ENH: add a float32-only device (#206)
* add an f32-only device, move device handling to _devices.py * refactor _info.py for device-specific dtypes * adapt creation functions for device-specific dtypes; add tests * manually sync dlpack enum updates to _devices.py As a part of a rebase/conflict resolution, this commit simply moves the dlpack/device updates to _devices.py. Co-authored-by: Tim Head <betatim@gmail.com> * MAINT: fix up after the rebase * add dlpack_device dunder to the F32_device * Update array_api_strict/_creation_functions.py Co-authored-by: Tim Head <betatim@gmail.com> * address review comments * raise if (device_type, device_id) is not found in the DLPack mappings * ENH: make device2 default to float32 (but still support float64) This way, "device2" mimics a pytorch CPU device (supports f64, defaults to f32), and "no_float64" mimics an MPS device (does not support double precision at all). While at it, fix the logic in asarray: whether a device does or does not support a dtype is different from what is the default dtype for this device. The the decision on the latter should not depend on the former. --------- Co-authored-by: Tim Head <betatim@gmail.com>
2 parents 28a48df + 572c4eb commit 34f301f

13 files changed

Lines changed: 505 additions & 182 deletions

‎array_api_strict/_array_object.py‎

Lines changed: 12 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from collections.abc import Iterator
2121
from enum import IntEnum
2222
from types import EllipsisType, ModuleType
23-
from typing import Any, Final, Literal, SupportsIndex, Callable
23+
from typing import Any, Literal, SupportsIndex, Callable
2424

2525
import numpy as np
2626
import numpy.typing as npt
@@ -40,64 +40,14 @@
4040
_real_to_complex_map,
4141
_result_type,
4242
)
43+
from ._devices import (
44+
CPU_DEVICE, Device, device_supports_dtype, _normalize_dl_device, _DLPACK_DEVICE_FOR,
45+
DLDeviceType
46+
)
4347
from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
4448
from ._typing import PyCapsule
4549

4650

47-
class Device:
48-
_device: Final[str]
49-
__slots__ = ("_device", "__weakref__")
50-
51-
def __init__(self, device: str = "CPU_DEVICE"):
52-
if device not in ("CPU_DEVICE", "device1", "device2"):
53-
raise ValueError(f"The device '{device}' is not a valid choice.")
54-
self._device = device
55-
56-
def __repr__(self) -> str:
57-
return f"array_api_strict.Device('{self._device}')"
58-
59-
def __eq__(self, other: object) -> bool:
60-
if not isinstance(other, Device):
61-
return False
62-
return self._device == other._device
63-
64-
def __hash__(self) -> int:
65-
return hash(("Device", self._device))
66-
67-
68-
CPU_DEVICE = Device()
69-
ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"))
70-
71-
72-
class DLDeviceType(IntEnum):
73-
kDLCPU = 1
74-
kDLCUDA = 2
75-
76-
77-
_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = {
78-
CPU_DEVICE: (DLDeviceType.kDLCPU, 0),
79-
Device("device1"): (DLDeviceType.kDLCUDA, 0),
80-
Device("device2"): (DLDeviceType.kDLCUDA, 1),
81-
}
82-
83-
_DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = {
84-
(int(device_type), device_id): logical_device
85-
for logical_device, (device_type, device_id) in _DLPACK_DEVICE_FOR.items()
86-
}
87-
88-
89-
def _normalize_dl_device(device_type: IntEnum | int, device_id: int) -> tuple[int, int]:
90-
return (int(device_type), device_id)
91-
92-
93-
def _device_from_dlpack_device(
94-
device_type: IntEnum | int, device_id: int
95-
) -> Device:
96-
return _DLPACK_DEVICE_TO_LOGICAL.get(
97-
_normalize_dl_device(device_type, device_id), CPU_DEVICE
98-
)
99-
100-
10151
class Array:
10252
"""
10353
n-d array object for the array API namespace.
@@ -142,10 +92,15 @@ def _new(cls, x: npt.NDArray[Any] | np.generic, /, device: Device | None) -> Arr
14292
raise TypeError(
14393
f"The array_api_strict namespace does not support the dtype '{x.dtype}'"
14494
)
145-
obj._array = x
146-
obj._dtype = _dtype
95+
14796
if device is None:
14897
device = CPU_DEVICE
98+
if not device_supports_dtype(device, _dtype):
99+
raise ValueError(f"Device {device!r} does not support dtype={_dtype!r}.")
100+
101+
obj._array = x
102+
obj._dtype = _dtype
103+
149104
obj._device = device
150105
return obj
151106

‎array_api_strict/_creation_functions.py‎

Lines changed: 86 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,11 @@
55

66
import numpy as np
77

8-
from ._dtypes import DType, _all_dtypes, _np_dtype
8+
from ._dtypes import DType, _all_dtypes, _np_dtype, bool as xp_bool
9+
from ._devices import (
10+
Device, device_supports_dtype, get_default_dtypes,
11+
check_device as _check_device
12+
)
913
from ._flags import get_array_api_strict_flags
1014
from ._typing import NestedSequence, SupportsBufferProtocol, SupportsDLPack
1115

@@ -14,7 +18,7 @@
1418
from typing_extensions import TypeIs
1519

1620
# Circular import
17-
from ._array_object import Array, Device
21+
from ._array_object import Array
1822

1923

2024
class Undef(Enum):
@@ -24,10 +28,15 @@ class Undef(Enum):
2428
_undef = Undef.UNDEF
2529

2630

27-
def _check_valid_dtype(dtype: DType | None) -> None:
31+
def _check_valid_dtype(dtype: DType | None, device: Device | None = None) -> None:
2832
# Note: Only spelling dtypes as the dtype objects is supported.
29-
if dtype not in (None,) + _all_dtypes:
30-
raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
33+
if dtype is not None:
34+
if dtype not in _all_dtypes:
35+
raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
36+
37+
if device is not None:
38+
if not device_supports_dtype(device, dtype):
39+
raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.")
3140

3241

3342
def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]:
@@ -38,18 +47,6 @@ def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]:
3847
return True
3948

4049

41-
def _check_device(device: Device | None) -> None:
42-
# _array_object imports in this file are inside the functions to avoid
43-
# circular imports
44-
from ._array_object import ALL_DEVICES, Device
45-
46-
if device is not None and not isinstance(device, Device):
47-
raise ValueError(f"Unsupported device {device!r}")
48-
49-
if device is not None and device not in ALL_DEVICES:
50-
raise ValueError(f"Unsupported device {device!r}")
51-
52-
5350
def asarray(
5451
obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol,
5552
/,
@@ -65,11 +62,12 @@ def asarray(
6562
"""
6663
from ._array_object import Array
6764

68-
_check_valid_dtype(dtype)
65+
_check_device(device)
66+
_check_valid_dtype(dtype, device)
6967
_np_dtype = None
7068
if dtype is not None:
7169
_np_dtype = dtype._np_dtype
72-
_check_device(device)
70+
7371
if isinstance(obj, Array) and device is None:
7472
device = obj.device
7573

@@ -108,6 +106,27 @@ def asarray(
108106
raise OverflowError("Integer out of bounds for array dtypes")
109107

110108
res = np.array(obj, dtype=_np_dtype, copy=copy)
109+
110+
# numpy default dtype may differ; if so, adjust the dtype
111+
if dtype is None and device is not None:
112+
res_dtype = DType(res.dtype)
113+
# The dtype selected by Numpy might not be the default dtype
114+
# on this device. We thus find the default dtype for the dtype "kind", and
115+
# cast to the device-appropriate default.
116+
from ._data_type_functions import isdtype
117+
if isdtype(res_dtype, "bool"):
118+
target_dtype = DType("bool")
119+
elif isdtype(res_dtype, "integral"):
120+
target_dtype = get_default_dtypes(device)["integral"]
121+
elif isdtype(res_dtype, "real floating"):
122+
target_dtype = get_default_dtypes(device)["real floating"]
123+
elif isdtype(res_dtype, "complex floating"):
124+
target_dtype = get_default_dtypes(device)["complex floating"]
125+
else:
126+
raise ValueError(f"{res_dtype = } not understood.")
127+
128+
res = res.astype(target_dtype._np_dtype)
129+
111130
return Array._new(res, device=device)
112131

113132

@@ -127,8 +146,13 @@ def arange(
127146
"""
128147
from ._array_object import Array
129148

130-
_check_valid_dtype(dtype)
131149
_check_device(device)
150+
_check_valid_dtype(dtype, device)
151+
if dtype is None:
152+
if any(isinstance(x, float) for x in (start, stop, step)):
153+
dtype = get_default_dtypes(device)["real floating"]
154+
else:
155+
dtype = get_default_dtypes(device)["integral"]
132156

133157
return Array._new(
134158
np.arange(start, stop, step, dtype=_np_dtype(dtype)),
@@ -149,8 +173,10 @@ def empty(
149173
"""
150174
from ._array_object import Array
151175

152-
_check_valid_dtype(dtype)
153176
_check_device(device)
177+
_check_valid_dtype(dtype, device)
178+
if dtype is None:
179+
dtype = get_default_dtypes(device)["real floating"]
154180

155181
return Array._new(np.empty(shape, dtype=_np_dtype(dtype)), device=device)
156182

@@ -165,10 +191,12 @@ def empty_like(
165191
"""
166192
from ._array_object import Array
167193

168-
_check_valid_dtype(dtype)
169194
_check_device(device)
170195
if device is None:
171196
device = x.device
197+
if dtype is None:
198+
dtype = x.dtype
199+
_check_valid_dtype(dtype, device)
172200

173201
return Array._new(np.empty_like(x._array, dtype=_np_dtype(dtype)), device=device)
174202

@@ -189,8 +217,10 @@ def eye(
189217
"""
190218
from ._array_object import Array
191219

192-
_check_valid_dtype(dtype)
193220
_check_device(device)
221+
_check_valid_dtype(dtype, device)
222+
if dtype is None:
223+
dtype = get_default_dtypes(device)["real floating"]
194224

195225
return Array._new(
196226
np.eye(n_rows, M=n_cols, k=k, dtype=_np_dtype(dtype)), device=device
@@ -217,7 +247,7 @@ def from_dlpack(
217247
else:
218248
device = None
219249
if hasattr(x, "__dlpack_device__"):
220-
from ._array_object import _device_from_dlpack_device
250+
from ._devices import _device_from_dlpack_device
221251

222252
dl_type, dl_id = x.__dlpack_device__()
223253
device = _device_from_dlpack_device(dl_type, dl_id)
@@ -242,12 +272,22 @@ def full(
242272
"""
243273
from ._array_object import Array
244274

245-
_check_valid_dtype(dtype)
246275
_check_device(device)
276+
_check_valid_dtype(dtype, device)
247277

248278
if not isinstance(fill_value, bool | int | float | complex):
249279
msg = f"Expected Python scalar fill_value, got type {type(fill_value)}"
250280
raise TypeError(msg)
281+
282+
if dtype is None:
283+
if type(fill_value) == bool:
284+
dtype = xp_bool
285+
else:
286+
kind = {
287+
int: "integral", float: "real floating", complex: "complex floating"
288+
}[type(fill_value)]
289+
dtype = get_default_dtypes(device)[kind]
290+
251291
res = np.full(shape, fill_value, dtype=_np_dtype(dtype))
252292
if DType(res.dtype) not in _all_dtypes:
253293
# This will happen if the fill value is not something that NumPy
@@ -271,10 +311,12 @@ def full_like(
271311
"""
272312
from ._array_object import Array
273313

274-
_check_valid_dtype(dtype)
275314
_check_device(device)
276315
if device is None:
277316
device = x.device
317+
if dtype is None:
318+
dtype = x.dtype
319+
_check_valid_dtype(dtype, device)
278320

279321
if not isinstance(fill_value, bool | int | float | complex):
280322
msg = f"Expected Python scalar fill_value, got type {type(fill_value)}"
@@ -305,8 +347,13 @@ def linspace(
305347
"""
306348
from ._array_object import Array
307349

308-
_check_valid_dtype(dtype)
309350
_check_device(device)
351+
_check_valid_dtype(dtype, device)
352+
if dtype is None:
353+
if isinstance(start, complex) or isinstance(stop, complex):
354+
dtype = get_default_dtypes(device)["complex floating"]
355+
else:
356+
dtype = get_default_dtypes(device)["real floating"]
310357

311358
return Array._new(
312359
np.linspace(start, stop, num, dtype=_np_dtype(dtype), endpoint=endpoint),
@@ -358,8 +405,10 @@ def ones(
358405
"""
359406
from ._array_object import Array
360407

361-
_check_valid_dtype(dtype)
362408
_check_device(device)
409+
_check_valid_dtype(dtype, device)
410+
if dtype is None:
411+
dtype = get_default_dtypes(device)["real floating"]
363412

364413
return Array._new(np.ones(shape, dtype=_np_dtype(dtype)), device=device)
365414

@@ -374,10 +423,12 @@ def ones_like(
374423
"""
375424
from ._array_object import Array
376425

377-
_check_valid_dtype(dtype)
378426
_check_device(device)
379427
if device is None:
380428
device = x.device
429+
if dtype is None:
430+
dtype = x.dtype
431+
_check_valid_dtype(dtype, device)
381432

382433
return Array._new(np.ones_like(x._array, dtype=_np_dtype(dtype)), device=device)
383434

@@ -423,8 +474,10 @@ def zeros(
423474
"""
424475
from ._array_object import Array
425476

426-
_check_valid_dtype(dtype)
427477
_check_device(device)
478+
_check_valid_dtype(dtype, device)
479+
if dtype is None:
480+
dtype = get_default_dtypes(device)["real floating"]
428481

429482
return Array._new(np.zeros(shape, dtype=_np_dtype(dtype)), device=device)
430483

@@ -439,9 +492,11 @@ def zeros_like(
439492
"""
440493
from ._array_object import Array
441494

442-
_check_valid_dtype(dtype)
443495
_check_device(device)
444496
if device is None:
445497
device = x.device
498+
if dtype is None:
499+
dtype = x.dtype
500+
_check_valid_dtype(dtype, device)
446501

447502
return Array._new(np.zeros_like(x._array, dtype=_np_dtype(dtype)), device=device)

0 commit comments

Comments
 (0)