From 2c489ca3ca4ce0bf217938ebd6e7aca78c0a59cb Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Tue, 30 Jun 2026 17:12:00 +0200 Subject: [PATCH 01/10] add an f32-only device, move device handling to _devices.py --- array_api_strict/_array_object.py | 66 +++---------------- array_api_strict/_devices.py | 101 ++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 57 deletions(-) create mode 100644 array_api_strict/_devices.py diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py index b356b82b..2caac0fd 100644 --- a/array_api_strict/_array_object.py +++ b/array_api_strict/_array_object.py @@ -20,7 +20,7 @@ from collections.abc import Iterator from enum import IntEnum from types import EllipsisType, ModuleType -from typing import Any, Final, Literal, SupportsIndex, Callable +from typing import Any, Literal, SupportsIndex, Callable import numpy as np import numpy.typing as npt @@ -40,64 +40,11 @@ _real_to_complex_map, _result_type, ) +from ._devices import CPU_DEVICE, Device, device_supports_dtype from ._flags import get_array_api_strict_flags, set_array_api_strict_flags from ._typing import PyCapsule -class Device: - _device: Final[str] - __slots__ = ("_device", "__weakref__") - - def __init__(self, device: str = "CPU_DEVICE"): - if device not in ("CPU_DEVICE", "device1", "device2"): - raise ValueError(f"The device '{device}' is not a valid choice.") - self._device = device - - def __repr__(self) -> str: - return f"array_api_strict.Device('{self._device}')" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, Device): - return False - return self._device == other._device - - def __hash__(self) -> int: - return hash(("Device", self._device)) - - -CPU_DEVICE = Device() -ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2")) - - -class DLDeviceType(IntEnum): - kDLCPU = 1 - kDLCUDA = 2 - - -_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = { - CPU_DEVICE: (DLDeviceType.kDLCPU, 0), - Device("device1"): (DLDeviceType.kDLCUDA, 0), - Device("device2"): (DLDeviceType.kDLCUDA, 1), -} - -_DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = { - (int(device_type), device_id): logical_device - for logical_device, (device_type, device_id) in _DLPACK_DEVICE_FOR.items() -} - - -def _normalize_dl_device(device_type: IntEnum | int, device_id: int) -> tuple[int, int]: - return (int(device_type), device_id) - - -def _device_from_dlpack_device( - device_type: IntEnum | int, device_id: int -) -> Device: - return _DLPACK_DEVICE_TO_LOGICAL.get( - _normalize_dl_device(device_type, device_id), CPU_DEVICE - ) - - class Array: """ n-d array object for the array API namespace. @@ -142,10 +89,15 @@ def _new(cls, x: npt.NDArray[Any] | np.generic, /, device: Device | None) -> Arr raise TypeError( f"The array_api_strict namespace does not support the dtype '{x.dtype}'" ) - obj._array = x - obj._dtype = _dtype + if device is None: device = CPU_DEVICE + if not device_supports_dtype(device, _dtype): + raise ValueError(f"Device {device!r} does not support dtype={_dtype!r}.") + + obj._array = x + obj._dtype = _dtype + obj._device = device return obj diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py new file mode 100644 index 00000000..cd76eae2 --- /dev/null +++ b/array_api_strict/_devices.py @@ -0,0 +1,101 @@ +from typing import Final + +from ._dtypes import ( + DType, float32, float64, complex64, complex128, int64, + _all_dtypes, _boolean_dtypes, _signed_integer_dtypes, + _unsigned_integer_dtypes, _integer_dtypes, _real_floating_dtypes, + _complex_floating_dtypes, _numeric_dtypes +) + +_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "F32_device") + +class Device: + _device: Final[str] + __slots__ = ("_device", "__weakref__") + + def __init__(self, device: str = "CPU_DEVICE"): + if device not in _ALL_DEVICE_NAMES: + raise ValueError(f"The device '{device}' is not a valid choice.") + self._device = device + + def __repr__(self) -> str: + return f"array_api_strict.Device('{self._device}')" + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Device): + return False + return self._device == other._device + + def __hash__(self) -> int: + return hash(("Device", self._device)) + + def _supported_dtypes(self) -> list[DType]: + # XXX useful? Unused ATM + return list(dt for dt in _all_dtypes if device_supports_dtype(self, dt)) + + +CPU_DEVICE = Device() +_F32_DEVICE = Device("F32_device") + +ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), _F32_DEVICE) + + +def check_device(device: Device | None) -> None: + if device is not None and not isinstance(device, Device): + raise ValueError(f"Unsupported device {device!r}") + + if device is not None and device not in ALL_DEVICES: + raise ValueError(f"Unsupported device {device!r}") + + +# Helpers for device-specific dtype support + +def get_default_dtypes(device: Device | None = None) -> dict[str, Device]: + if device == _F32_DEVICE: + return { + "real floating": float32, + "complex floating": complex64, + "integral": int64, + "indexing": int64, + } + else: + return { + "real floating": float64, + "complex floating": complex128, + "integral": int64, + "indexing": int64, + } + + +def device_supports_dtype(device: Device | None, dtype: DType |None) -> bool: + """True if `device` supports `dtype`, False otherwise.""" + # special-case F32_device + if device == _F32_DEVICE: + return dtype not in (float64, complex128) + + # All other devices support all dtypes + return True + + +def _map_supported(dtypes: list[DType], device: Device) -> dict[str, DType]: + return { + dt._canonic_name: dt + for dt in dtypes + if device_supports_dtype(device, dt) + } + + +# _info.dtypes() maps "kind" -> dict of {name: dtype} +# Note that "kinds" differ from "categories" above, per the spec. + +_kind_to_dtypes = { + None: _all_dtypes, + "bool": _boolean_dtypes, + "signed integer": _signed_integer_dtypes, + "unsigned integer": _unsigned_integer_dtypes, + "integral": _integer_dtypes, + "real floating": _real_floating_dtypes, + "complex floating": _complex_floating_dtypes, + "numeric": _numeric_dtypes +} + From f321cdc40b4e4550d6481bf155d0a94c40661702 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Tue, 30 Jun 2026 17:13:27 +0200 Subject: [PATCH 02/10] refactor _info.py for device-specific dtypes --- array_api_strict/_info.py | 94 +++++++-------------------------------- 1 file changed, 16 insertions(+), 78 deletions(-) diff --git a/array_api_strict/_info.py b/array_api_strict/_info.py index 12beed0f..cbee0360 100644 --- a/array_api_strict/_info.py +++ b/array_api_strict/_info.py @@ -1,7 +1,7 @@ import numpy as np -from . import _dtypes as dt -from ._array_object import ALL_DEVICES, CPU_DEVICE, Device +from . import _devices +from ._devices import ALL_DEVICES, CPU_DEVICE, Device from ._flags import get_array_api_strict_flags, requires_api_version from ._typing import Capabilities, DataTypes, DefaultDataTypes @@ -40,12 +40,7 @@ def default_dtypes( *, device: Device | None = None, ) -> DefaultDataTypes: - return { - "real floating": dt.float64, - "complex floating": dt.complex128, - "integral": dt.int64, - "indexing": dt.int64, - } + return _devices.get_default_dtypes(device) @requires_api_version('2023.12') def dtypes( @@ -54,78 +49,21 @@ def dtypes( device: Device | None = None, kind: str | tuple[str, ...] | None = None, ) -> DataTypes: - if kind is None: - return { - "bool": dt.bool, - "int8": dt.int8, - "int16": dt.int16, - "int32": dt.int32, - "int64": dt.int64, - "uint8": dt.uint8, - "uint16": dt.uint16, - "uint32": dt.uint32, - "uint64": dt.uint64, - "float32": dt.float32, - "float64": dt.float64, - "complex64": dt.complex64, - "complex128": dt.complex128, - } - if kind == "bool": - return {"bool": dt.bool} - if kind == "signed integer": - return { - "int8": dt.int8, - "int16": dt.int16, - "int32": dt.int32, - "int64": dt.int64, - } - if kind == "unsigned integer": - return { - "uint8": dt.uint8, - "uint16": dt.uint16, - "uint32": dt.uint32, - "uint64": dt.uint64, - } - if kind == "integral": - return { - "int8": dt.int8, - "int16": dt.int16, - "int32": dt.int32, - "int64": dt.int64, - "uint8": dt.uint8, - "uint16": dt.uint16, - "uint32": dt.uint32, - "uint64": dt.uint64, - } - if kind == "real floating": - return { - "float32": dt.float32, - "float64": dt.float64, - } - if kind == "complex floating": - return { - "complex64": dt.complex64, - "complex128": dt.complex128, - } - if kind == "numeric": - return { - "int8": dt.int8, - "int16": dt.int16, - "int32": dt.int32, - "int64": dt.int64, - "uint8": dt.uint8, - "uint16": dt.uint16, - "uint32": dt.uint32, - "uint64": dt.uint64, - "float32": dt.float32, - "float64": dt.float64, - "complex64": dt.complex64, - "complex128": dt.complex128, - } - if isinstance(kind, tuple): + if device is None: + device = CPU_DEVICE + if isinstance(kind, type(None) | str): + + try: + dtypes = _devices._kind_to_dtypes[kind] + except KeyError: + raise ValueError(f"unsupported kind: {kind!r}") + res = _devices._map_supported(dtypes, device) + return res + + elif isinstance(kind, tuple): res: DataTypes = {} for k in kind: - res.update(self.dtypes(kind=k)) + res.update(self.dtypes(kind=k, device=device)) return res raise ValueError(f"unsupported kind: {kind!r}") From 94786bf2dc88d64c508e809ebcaf0d5a2c56403d Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Tue, 30 Jun 2026 17:15:09 +0200 Subject: [PATCH 03/10] adapt creation functions for device-specific dtypes; add tests --- array_api_strict/_creation_functions.py | 115 +++++++++++---- array_api_strict/_dtypes.py | 4 +- array_api_strict/_fft.py | 22 ++- array_api_strict/_statistical_functions.py | 10 +- .../tests/test_creation_functions.py | 139 +++++++++++++++++- array_api_strict/tests/test_device_support.py | 56 ++++++- .../tests/test_elementwise_functions.py | 7 +- .../tests/test_searching_functions.py | 2 +- .../tests/test_statistical_functions.py | 16 ++ 9 files changed, 326 insertions(+), 45 deletions(-) diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index a5320284..b7f7b6e3 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -5,7 +5,11 @@ import numpy as np -from ._dtypes import DType, _all_dtypes, _np_dtype +from ._dtypes import DType, _all_dtypes, _np_dtype, bool as xp_bool +from ._devices import ( + Device, device_supports_dtype, get_default_dtypes, + check_device as _check_device +) from ._flags import get_array_api_strict_flags from ._typing import NestedSequence, SupportsBufferProtocol, SupportsDLPack @@ -14,7 +18,7 @@ from typing_extensions import TypeIs # Circular import - from ._array_object import Array, Device + from ._array_object import Array class Undef(Enum): @@ -24,10 +28,15 @@ class Undef(Enum): _undef = Undef.UNDEF -def _check_valid_dtype(dtype: DType | None) -> None: +def _check_valid_dtype(dtype: DType | None, device: Device | None = None) -> None: # Note: Only spelling dtypes as the dtype objects is supported. - if dtype not in (None,) + _all_dtypes: - raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}") + if dtype is not None: + if dtype not in _all_dtypes: + raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}") + + if device is not None: + if not device_supports_dtype(device, dtype): + raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.") def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]: @@ -38,18 +47,6 @@ def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]: return True -def _check_device(device: Device | None) -> None: - # _array_object imports in this file are inside the functions to avoid - # circular imports - from ._array_object import ALL_DEVICES, Device - - if device is not None and not isinstance(device, Device): - raise ValueError(f"Unsupported device {device!r}") - - if device is not None and device not in ALL_DEVICES: - raise ValueError(f"Unsupported device {device!r}") - - def asarray( obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol, /, @@ -65,11 +62,12 @@ def asarray( """ from ._array_object import Array - _check_valid_dtype(dtype) + _check_device(device) + _check_valid_dtype(dtype, device) _np_dtype = None if dtype is not None: _np_dtype = dtype._np_dtype - _check_device(device) + if isinstance(obj, Array) and device is None: device = obj.device @@ -108,6 +106,27 @@ def asarray( raise OverflowError("Integer out of bounds for array dtypes") res = np.array(obj, dtype=_np_dtype, copy=copy) + + # numpy default dtype may differ; if so, adjust the dtype + if dtype is None and device is not None: + res_dtype = DType(res.dtype) + if not device_supports_dtype(device, res_dtype): + # find out the default dtype for the device + from ._data_type_functions import isdtype + if isdtype(res_dtype, "bool"): + targ_dtype = DType("bool") + elif isdtype(res_dtype, "integral"): + targ_dtype = get_default_dtypes(device)["integral"] + elif isdtype(res_dtype, "real floating"): + targ_dtype = get_default_dtypes(device)["real floating"] + elif isdtype(res_dtype, "complex floating"): + targ_dtype = get_default_dtypes(device)["complex floating"] + else: + raise ValueError(f"{res_dtype = } not understood.") + del isdtype + + res = res.astype(targ_dtype._np_dtype) + return Array._new(res, device=device) @@ -127,8 +146,13 @@ def arange( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) + _check_valid_dtype(dtype, device) + if dtype is None: + if any(isinstance(x, float) for x in (start, stop, step)): + dtype = get_default_dtypes(device)["real floating"] + else: + dtype = get_default_dtypes(device)["integral"] return Array._new( np.arange(start, stop, step, dtype=_np_dtype(dtype)), @@ -149,8 +173,10 @@ def empty( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) + _check_valid_dtype(dtype, device) + if dtype is None: + dtype = get_default_dtypes(device)["real floating"] return Array._new(np.empty(shape, dtype=_np_dtype(dtype)), device=device) @@ -165,10 +191,12 @@ def empty_like( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) if device is None: device = x.device + if dtype is None: + dtype = x.dtype + _check_valid_dtype(dtype, device) return Array._new(np.empty_like(x._array, dtype=_np_dtype(dtype)), device=device) @@ -189,8 +217,10 @@ def eye( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) + _check_valid_dtype(dtype, device) + if dtype is None: + dtype = get_default_dtypes(device)["real floating"] return Array._new( np.eye(n_rows, M=n_cols, k=k, dtype=_np_dtype(dtype)), device=device @@ -242,12 +272,22 @@ def full( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) + _check_valid_dtype(dtype, device) if not isinstance(fill_value, bool | int | float | complex): msg = f"Expected Python scalar fill_value, got type {type(fill_value)}" raise TypeError(msg) + + if dtype is None: + if type(fill_value) == bool: + dtype = xp_bool + else: + kind = { + int: "integral", float: "real floating", complex: "complex floating" + }[type(fill_value)] + dtype = get_default_dtypes(device)[kind] + res = np.full(shape, fill_value, dtype=_np_dtype(dtype)) if DType(res.dtype) not in _all_dtypes: # This will happen if the fill value is not something that NumPy @@ -271,10 +311,12 @@ def full_like( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) if device is None: device = x.device + if dtype is None: + dtype = x.dtype + _check_valid_dtype(dtype, device) if not isinstance(fill_value, bool | int | float | complex): msg = f"Expected Python scalar fill_value, got type {type(fill_value)}" @@ -305,8 +347,13 @@ def linspace( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) + _check_valid_dtype(dtype, device) + if dtype is None: + if isinstance(start, complex) or isinstance(stop, complex): + dtype = get_default_dtypes(device)["complex floating"] + else: + dtype = get_default_dtypes(device)["real floating"] return Array._new( np.linspace(start, stop, num, dtype=_np_dtype(dtype), endpoint=endpoint), @@ -358,8 +405,10 @@ def ones( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) + _check_valid_dtype(dtype, device) + if dtype is None: + dtype = get_default_dtypes(device)["real floating"] return Array._new(np.ones(shape, dtype=_np_dtype(dtype)), device=device) @@ -374,10 +423,12 @@ def ones_like( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) if device is None: device = x.device + if dtype is None: + dtype = x.dtype + _check_valid_dtype(dtype, device) return Array._new(np.ones_like(x._array, dtype=_np_dtype(dtype)), device=device) @@ -423,8 +474,10 @@ def zeros( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) + _check_valid_dtype(dtype, device) + if dtype is None: + dtype = get_default_dtypes(device)["real floating"] return Array._new(np.zeros(shape, dtype=_np_dtype(dtype)), device=device) @@ -439,9 +492,11 @@ def zeros_like( """ from ._array_object import Array - _check_valid_dtype(dtype) _check_device(device) if device is None: device = x.device + if dtype is None: + dtype = x.dtype + _check_valid_dtype(dtype, device) return Array._new(np.zeros_like(x._array, dtype=_np_dtype(dtype)), device=device) diff --git a/array_api_strict/_dtypes.py b/array_api_strict/_dtypes.py index 564db5a8..2d5ec33f 100644 --- a/array_api_strict/_dtypes.py +++ b/array_api_strict/_dtypes.py @@ -11,9 +11,11 @@ class DType: _np_dtype: Final[np.dtype[Any]] - __slots__ = ("_np_dtype", "__weakref__") + _canonic_name: Final[Any] + __slots__ = ("_np_dtype", "_canonic_name", "__weakref__") def __init__(self, np_dtype: npt.DTypeLike): + self._canonic_name = np_dtype self._np_dtype = np.dtype(np_dtype) def __repr__(self) -> str: diff --git a/array_api_strict/_fft.py b/array_api_strict/_fft.py index c2c617e1..e069435f 100644 --- a/array_api_strict/_fft.py +++ b/array_api_strict/_fft.py @@ -3,7 +3,8 @@ import numpy as np -from ._array_object import ALL_DEVICES, Array, Device +from ._array_object import Array +from ._devices import ALL_DEVICES, Device, device_supports_dtype from ._data_type_functions import astype from ._dtypes import ( DType, @@ -13,6 +14,7 @@ complex64, float32, ) +from ._info import __array_namespace_info__ from ._flags import requires_extension @@ -268,6 +270,15 @@ def fftfreq( np_result = np.fft.fftfreq(n, d=d) if dtype: np_result = np_result.astype(dtype._np_dtype) + + if not device_supports_dtype(device, DType(np_result.dtype)): + if dtype: + # user input unsupported + raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.") + + dt = __array_namespace_info__().default_dtypes(device=device)["real floating"] + np_result = np_result.astype(dt._np_dtype) + return Array._new(np_result, device=device) @requires_extension('fft') @@ -292,6 +303,15 @@ def rfftfreq( np_result = np.fft.rfftfreq(n, d=d) if dtype: np_result = np_result.astype(dtype._np_dtype) + + if not device_supports_dtype(device, DType(np_result.dtype)): + if dtype: + # user input unsupported + raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.") + + dt = __array_namespace_info__().default_dtypes(device=device)["real floating"] + np_result = np_result.astype(dt._np_dtype) + return Array._new(np_result, device=device) @requires_extension('fft') diff --git a/array_api_strict/_statistical_functions.py b/array_api_strict/_statistical_functions.py index 35876dda..bfcbe0c7 100644 --- a/array_api_strict/_statistical_functions.py +++ b/array_api_strict/_statistical_functions.py @@ -39,7 +39,10 @@ def cumulative_sum( if include_initial: if axis < 0: axis += x.ndim - x = concat([zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis) + x = concat( + [zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype, device=x.device), x], + axis=axis + ) return Array._new(np.cumsum(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device) @@ -66,7 +69,10 @@ def cumulative_prod( if include_initial: if axis < 0: axis += x.ndim - x = concat([ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis) + x = concat( + [ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype, device=x.device), x], + axis=axis + ) return Array._new(np.cumprod(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device) diff --git a/array_api_strict/tests/test_creation_functions.py b/array_api_strict/tests/test_creation_functions.py index 79202893..ee1263ae 100644 --- a/array_api_strict/tests/test_creation_functions.py +++ b/array_api_strict/tests/test_creation_functions.py @@ -22,8 +22,10 @@ zeros, zeros_like, ) -from .._dtypes import float32, float64 -from .._array_object import Array, CPU_DEVICE, Device +from .._dtypes import float32, float64, bool as xp_bool +from .._array_object import Array +from .._devices import CPU_DEVICE, ALL_DEVICES, Device +from .._info import __array_namespace_info__ from .._flags import set_array_api_strict_flags def test_asarray_errors(): @@ -212,6 +214,7 @@ def test_zeros_like_errors(): assert_raises(ValueError, lambda: zeros_like(asarray(1), dtype=int)) assert_raises(ValueError, lambda: zeros_like(asarray(1), dtype="i")) + def test_meshgrid_dtype_errors(): # Doesn't raise meshgrid() @@ -221,6 +224,138 @@ def test_meshgrid_dtype_errors(): assert_raises(ValueError, lambda: meshgrid(asarray([1.], dtype=float32), asarray([1.], dtype=float64))) + +def _full(a, *args, **kwds): + return full(a, fill_value=42.0, *args, **kwds) + + +def _full_like(a, *args, **kwds): + return full_like(a, fill_value=42.0, *args, **kwds) + + +class TestDefaultDType: + + info = __array_namespace_info__() + + @pytest.mark.parametrize("device", ALL_DEVICES) + @pytest.mark.parametrize("func", [empty, zeros, ones, _full]) + def test_ones_etc(self, func, device): + a = func(1, device=device) + assert a.dtype == self.info.default_dtypes(device=device)["real floating"] + + @pytest.mark.parametrize("func", [empty_like, zeros_like, ones_like, _full_like]) + def test_ones_like_etc_correct(self, func): + # float32 is preserved + a = ones(2, dtype=float32) + device = Device('F32_device') + b = func(a, device=device) + assert b.dtype == self.info.default_dtypes(device=device)["real floating"] + assert b.device == device + + @pytest.mark.parametrize("func", [empty_like, zeros_like, ones_like, _full_like]) + def test_ones_like_etc_incorrect(self, func): + a = ones(2) + assert a.dtype == float64 + assert a.device == Device() + + # XXX: a.dtype not supported by the device: ValueError or TypeError? + + # >>> a = torch.ones(3, dtype=torch.float64, device='cpu') + # >>> torch.ones_like(a, device='mps') + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework + # doesn't support float64. + + # incompatible dtype inferred from `a.dtype` + with pytest.raises((TypeError, ValueError)): + func(a, device=Device('F32_device')) + + # `a.dtype` is compatible but the explicit dtype= argument is incompatible + a = ones(2, dtype=float32) + with pytest.raises((TypeError, ValueError)): + func(a, device=Device('F32_device'), dtype=float64) + + def test_eye(self): + device = Device('F32_device') + a = eye(3, device=device) + assert a.dtype == self.info.default_dtypes(device=device)["real floating"] + assert a.device == device + + with pytest.raises((TypeError, ValueError)): + eye(3, device=device, dtype=float64) + + def test_linspace(self): + device = Device('F32_device') + + a = linspace(1, 10, 11, device=device) + assert a.dtype == self.info.default_dtypes(device=device)["real floating"] + assert a.device == device + + a = linspace(1+0j, 10, 11, device=device) + assert a.dtype == self.info.default_dtypes(device=device)["complex floating"] + + with pytest.raises((TypeError, ValueError)): + linspace(1, 10, 11, device=device, dtype=float64) + + def test_arange(self): + device = Device('F32_device') + + a = arange(0, 10, 1, device=device) + assert a.dtype == self.info.default_dtypes(device=device)["integral"] + assert a.device == device + + a = arange(0.0, 10, 1, device=device) + assert a.dtype == self.info.default_dtypes(device=device)["real floating"] + assert a.device == device + + with pytest.raises((TypeError, ValueError)): + arange(0, 10, 1, device=device, dtype=float64) + + with pytest.raises((TypeError, ValueError)): + arange(0.0, 10, 1, device=device, dtype=float64) + + def test_asarray(self): + device = Device('F32_device') + + ### asarray(python_object) + for x in (True, [False,]): + arr = asarray(x, device=device) + assert arr.dtype == xp_bool + assert arr.device == device + + for x in [1, [1,]]: + arr = asarray(x, device=device) + assert arr.dtype == self.info.default_dtypes(device=device)['integral'] + assert arr.device == device + + for x in [1.0, [1.0,]]: + arr = asarray(x, device=device) + assert arr.dtype == self.info.default_dtypes(device=device)['real floating'] + assert arr.device == device + + for x in [1j, [1j,]]: + arr = asarray(x, device=device) + assert arr.dtype == self.info.default_dtypes(device=device)['complex floating'] + assert arr.device == device + + # asarray(python_object, dtype=unsupported_by_device) + with pytest.raises(ValueError, match="Device"): + asarray(1, dtype=float64, device=device) + + ### asarray(array) + + # compatible dtypes, device transfer + src = asarray(1, dtype=float32, device=Device('device1')) + dst = asarray(src, device=device) + assert dst.device == device + assert dst.dtype == float32 + + # incompatible dtypes, device transfer + src = asarray(1, dtype=float64, device=Device('device1')) + + with pytest.raises(ValueError, match="Device"): + asarray(src, device=device) + + @pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12']) def from_dlpack_2023_12(api_version): if api_version != '2022.12': diff --git a/array_api_strict/tests/test_device_support.py b/array_api_strict/tests/test_device_support.py index 0f3d6b50..84954db8 100644 --- a/array_api_strict/tests/test_device_support.py +++ b/array_api_strict/tests/test_device_support.py @@ -1,6 +1,6 @@ import pytest -import array_api_strict +import array_api_strict as xp @pytest.mark.parametrize( @@ -18,11 +18,11 @@ ), ) def test_fft_device_support_complex(func_name): - func = getattr(array_api_strict.fft, func_name) - x = array_api_strict.asarray( + func = getattr(xp.fft, func_name) + x = xp.asarray( [1, 2.0], - dtype=array_api_strict.complex64, - device=array_api_strict.Device("device1"), + dtype=xp.complex64, + device=xp.Device("device1"), ) y = func(x) @@ -31,8 +31,50 @@ def test_fft_device_support_complex(func_name): @pytest.mark.parametrize("func_name", ("rfft", "rfftn", "ihfft")) def test_fft_device_support_real(func_name): - func = getattr(array_api_strict.fft, func_name) - x = array_api_strict.asarray([1, 2.0], device=array_api_strict.Device("device1")) + func = getattr(xp.fft, func_name) + x = xp.asarray([1, 2.0], device=xp.Device("device1")) y = func(x) assert x.device == y.device + + +@pytest.mark.parametrize("func_name", ("fftfreq", "rfftfreq")) +def test_fft_default_dtype(func_name): + func = getattr(xp.fft, func_name) + device = xp.Device("F32_device") + res = func(3, device=device) + assert res.device == device + assert res.dtype == xp.__array_namespace_info__().default_dtypes(device=device)["real floating"] + + with pytest.raises(ValueError): + func(3, device=device, dtype=xp.float64) + + +class TestF32Device: + @pytest.mark.parametrize("dtype_str", ["float64", "complex128"]) + def test_f64_raises(self, dtype_str): + f32_device = xp.Device("F32_device") + dtype = getattr(xp, dtype_str) + with pytest.raises(ValueError): + xp.arange(3, device=f32_device, dtype=dtype) + + def test_info_no_f64(self): + f32_device = xp.Device("F32_device") + + info = xp.__array_namespace_info__() + all_dtypes = info.dtypes(device=f32_device) + assert "float64" not in all_dtypes + assert "complex128" not in all_dtypes + + def test_info_default_dtypes(self): + f32_device = xp.Device("F32_device") + info = xp.__array_namespace_info__() + defaults = info.default_dtypes(device=f32_device) + assert defaults["real floating"] == xp.float32 + assert defaults["complex floating"] == xp.complex64 + + cpu_device = xp.Device() + info = xp.__array_namespace_info__() + defaults = info.default_dtypes(device=cpu_device) + assert defaults["real floating"] == xp.float64 + assert defaults["complex floating"] == xp.complex128 diff --git a/array_api_strict/tests/test_elementwise_functions.py b/array_api_strict/tests/test_elementwise_functions.py index 050f2bcd..7fb6e339 100644 --- a/array_api_strict/tests/test_elementwise_functions.py +++ b/array_api_strict/tests/test_elementwise_functions.py @@ -6,7 +6,7 @@ from .. import asarray, _elementwise_functions -from .._array_object import ALL_DEVICES, CPU_DEVICE, Device +from .._devices import ALL_DEVICES, CPU_DEVICE, Device from .._elementwise_functions import bitwise_left_shift, bitwise_right_shift from .._dtypes import ( _dtype_categories, @@ -21,6 +21,7 @@ int64, uint64, ) +from .._info import __array_namespace_info__ from .test_array_object import _check_op_array_scalar, BIG_INT import array_api_strict @@ -144,6 +145,10 @@ def _array_vals(dtypes): yield asarray(1., dtype=dtype, device=device) dtypes = _dtype_categories[types] + + supported_dtypes = __array_namespace_info__().dtypes(device=device) + dtypes = [dt for dt in dtypes if dt in supported_dtypes] + func = getattr(_elementwise_functions, func_name) for x in _array_vals(dtypes): diff --git a/array_api_strict/tests/test_searching_functions.py b/array_api_strict/tests/test_searching_functions.py index abe1949c..409deb5b 100644 --- a/array_api_strict/tests/test_searching_functions.py +++ b/array_api_strict/tests/test_searching_functions.py @@ -3,7 +3,7 @@ import array_api_strict as xp from array_api_strict import ArrayAPIStrictFlags -from .._array_object import ALL_DEVICES, CPU_DEVICE, Device +from .._devices import ALL_DEVICES, CPU_DEVICE, Device from .._dtypes import _all_dtypes diff --git a/array_api_strict/tests/test_statistical_functions.py b/array_api_strict/tests/test_statistical_functions.py index d702b17b..52ad53a0 100644 --- a/array_api_strict/tests/test_statistical_functions.py +++ b/array_api_strict/tests/test_statistical_functions.py @@ -55,3 +55,19 @@ def test_mean_complex(): with pytest.raises(TypeError): xp.mean(xp.arange(3)) + +def test_cumsum_device(): + x = xp.arange(3, device=xp.Device('device1')) + y = xp.cumulative_sum(x, include_initial=True) + expected = xp.asarray([0, 0, 1, 3], device=x.device) + assert y.device == expected.device + assert xp.all(y == expected) + + +def test_cumprod_device(): + x = xp.arange(1, 4, device=xp.Device('device1')) + y = xp.cumulative_prod(x, include_initial=True) + expected = xp.asarray([1, 1, 2, 6], device=x.device) + assert y.device == expected.device + assert xp.all(y == expected) + From 083212060b10d00fda24553937ae97d9e96e4062 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Tue, 30 Jun 2026 17:25:48 +0200 Subject: [PATCH 04/10] 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 --- array_api_strict/_devices.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index cd76eae2..81df6d8d 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -1,4 +1,5 @@ from typing import Final +from enum import IntEnum from ._dtypes import ( DType, float32, float64, complex64, complex128, int64, @@ -39,6 +40,34 @@ def _supported_dtypes(self) -> list[DType]: ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), _F32_DEVICE) +class DLDeviceType(IntEnum): + kDLCPU = 1 + kDLCUDA = 2 + + +_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = { + CPU_DEVICE: (DLDeviceType.kDLCPU, 0), + Device("device1"): (DLDeviceType.kDLCUDA, 0), + Device("device2"): (DLDeviceType.kDLCUDA, 1), +} + +_DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = { + (int(device_type), device_id): logical_device + for logical_device, (device_type, device_id) in _DLPACK_DEVICE_FOR.items() +} + + +def _normalize_dl_device(device_type: IntEnum | int, device_id: int) -> tuple[int, int]: + return (int(device_type), device_id) + + +def _device_from_dlpack_device( + device_type: IntEnum | int, device_id: int +) -> Device: + return _DLPACK_DEVICE_TO_LOGICAL.get( + _normalize_dl_device(device_type, device_id), CPU_DEVICE + ) + def check_device(device: Device | None) -> None: if device is not None and not isinstance(device, Device): From 346c59a19d148a44490073fea46cb3d4209d3d10 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Tue, 30 Jun 2026 17:32:42 +0200 Subject: [PATCH 05/10] MAINT: fix up after the rebase --- array_api_strict/_array_object.py | 5 ++++- array_api_strict/_creation_functions.py | 2 +- array_api_strict/tests/test_array_object.py | 3 ++- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py index 2caac0fd..0587ec51 100644 --- a/array_api_strict/_array_object.py +++ b/array_api_strict/_array_object.py @@ -40,7 +40,10 @@ _real_to_complex_map, _result_type, ) -from ._devices import CPU_DEVICE, Device, device_supports_dtype +from ._devices import ( + CPU_DEVICE, Device, device_supports_dtype, _normalize_dl_device, _DLPACK_DEVICE_FOR, + DLDeviceType +) from ._flags import get_array_api_strict_flags, set_array_api_strict_flags from ._typing import PyCapsule diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index b7f7b6e3..92f9869a 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -247,7 +247,7 @@ def from_dlpack( else: device = None if hasattr(x, "__dlpack_device__"): - from ._array_object import _device_from_dlpack_device + from ._devices import _device_from_dlpack_device dl_type, dl_id = x.__dlpack_device__() device = _device_from_dlpack_device(dl_type, dl_id) diff --git a/array_api_strict/tests/test_array_object.py b/array_api_strict/tests/test_array_object.py index 5a50d8bf..5464e8e5 100644 --- a/array_api_strict/tests/test_array_object.py +++ b/array_api_strict/tests/test_array_object.py @@ -9,7 +9,8 @@ import pytest from .. import ones, arange, reshape, asarray, result_type, all, equal, stack -from .._array_object import Array, CPU_DEVICE, Device, DLDeviceType +from .._array_object import Array +from .._devices import CPU_DEVICE, Device, DLDeviceType from .._dtypes import ( _all_dtypes, _boolean_dtypes, From dbe38b85178bb62c6b21ea2377f1424d8a9f1c7d Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Sat, 4 Jul 2026 09:59:36 +0200 Subject: [PATCH 06/10] add dlpack_device dunder to the F32_device --- array_api_strict/_devices.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index 81df6d8d..a06b4344 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -43,12 +43,14 @@ def _supported_dtypes(self) -> list[DType]: class DLDeviceType(IntEnum): kDLCPU = 1 kDLCUDA = 2 + kDLMETAL = 8 _DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = { CPU_DEVICE: (DLDeviceType.kDLCPU, 0), Device("device1"): (DLDeviceType.kDLCUDA, 0), Device("device2"): (DLDeviceType.kDLCUDA, 1), + _F32_DEVICE: (DLDeviceType.kDLMETAL, 0), } _DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = { From 826f8f27589da25c4a7a2c65e90b422126b1ba75 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Mon, 6 Jul 2026 16:45:26 +0200 Subject: [PATCH 07/10] Update array_api_strict/_creation_functions.py Co-authored-by: Tim Head --- array_api_strict/_creation_functions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index 92f9869a..6bd1c336 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -111,7 +111,9 @@ def asarray( if dtype is None and device is not None: res_dtype = DType(res.dtype) if not device_supports_dtype(device, res_dtype): - # find out the default dtype for the device + # The dtype selected by Numpy might not be the default dtype + # on this device. If the dtype is not supported by the device we + # try to find one that is, aka the default dtype of this device. from ._data_type_functions import isdtype if isdtype(res_dtype, "bool"): targ_dtype = DType("bool") From df1574d03b7f73125f340d771718e642c10a5627 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Mon, 6 Jul 2026 16:45:47 +0200 Subject: [PATCH 08/10] address review comments --- array_api_strict/_creation_functions.py | 11 +++++----- array_api_strict/_devices.py | 22 ++++++++----------- array_api_strict/_dtypes.py | 6 ++--- .../tests/test_creation_functions.py | 14 ++++++------ array_api_strict/tests/test_device_support.py | 14 ++++++------ 5 files changed, 31 insertions(+), 36 deletions(-) diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index 6bd1c336..52a0be47 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -116,18 +116,17 @@ def asarray( # try to find one that is, aka the default dtype of this device. from ._data_type_functions import isdtype if isdtype(res_dtype, "bool"): - targ_dtype = DType("bool") + target_dtype = DType("bool") elif isdtype(res_dtype, "integral"): - targ_dtype = get_default_dtypes(device)["integral"] + target_dtype = get_default_dtypes(device)["integral"] elif isdtype(res_dtype, "real floating"): - targ_dtype = get_default_dtypes(device)["real floating"] + target_dtype = get_default_dtypes(device)["real floating"] elif isdtype(res_dtype, "complex floating"): - targ_dtype = get_default_dtypes(device)["complex floating"] + target_dtype = get_default_dtypes(device)["complex floating"] else: raise ValueError(f"{res_dtype = } not understood.") - del isdtype - res = res.astype(targ_dtype._np_dtype) + res = res.astype(target_dtype._np_dtype) return Array._new(res, device=device) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index a06b4344..d23dcba3 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -8,7 +8,7 @@ _complex_floating_dtypes, _numeric_dtypes ) -_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "F32_device") +_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "no_float64") class Device: _device: Final[str] @@ -30,15 +30,11 @@ def __eq__(self, other: object) -> bool: def __hash__(self) -> int: return hash(("Device", self._device)) - def _supported_dtypes(self) -> list[DType]: - # XXX useful? Unused ATM - return list(dt for dt in _all_dtypes if device_supports_dtype(self, dt)) - CPU_DEVICE = Device() -_F32_DEVICE = Device("F32_device") +NO_FLOAT64_DEVICE = Device("no_float64") -ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), _F32_DEVICE) +ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), NO_FLOAT64_DEVICE) class DLDeviceType(IntEnum): kDLCPU = 1 @@ -50,7 +46,7 @@ class DLDeviceType(IntEnum): CPU_DEVICE: (DLDeviceType.kDLCPU, 0), Device("device1"): (DLDeviceType.kDLCUDA, 0), Device("device2"): (DLDeviceType.kDLCUDA, 1), - _F32_DEVICE: (DLDeviceType.kDLMETAL, 0), + NO_FLOAT64_DEVICE: (DLDeviceType.kDLMETAL, 0), } _DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = { @@ -81,8 +77,8 @@ def check_device(device: Device | None) -> None: # Helpers for device-specific dtype support -def get_default_dtypes(device: Device | None = None) -> dict[str, Device]: - if device == _F32_DEVICE: +def get_default_dtypes(device: Device | None = None) -> dict[str, DType]: + if device == NO_FLOAT64_DEVICE: return { "real floating": float32, "complex floating": complex64, @@ -100,8 +96,8 @@ def get_default_dtypes(device: Device | None = None) -> dict[str, Device]: def device_supports_dtype(device: Device | None, dtype: DType |None) -> bool: """True if `device` supports `dtype`, False otherwise.""" - # special-case F32_device - if device == _F32_DEVICE: + # Device("no_float64") supports all dtypes except float64 and complex128 + if device == NO_FLOAT64_DEVICE: return dtype not in (float64, complex128) # All other devices support all dtypes @@ -110,7 +106,7 @@ def device_supports_dtype(device: Device | None, dtype: DType |None) -> bool: def _map_supported(dtypes: list[DType], device: Device) -> dict[str, DType]: return { - dt._canonic_name: dt + dt._canonical_name: dt for dt in dtypes if device_supports_dtype(device, dt) } diff --git a/array_api_strict/_dtypes.py b/array_api_strict/_dtypes.py index 2d5ec33f..158d8d96 100644 --- a/array_api_strict/_dtypes.py +++ b/array_api_strict/_dtypes.py @@ -11,11 +11,11 @@ class DType: _np_dtype: Final[np.dtype[Any]] - _canonic_name: Final[Any] - __slots__ = ("_np_dtype", "_canonic_name", "__weakref__") + _canonical_name: Final[Any] + __slots__ = ("_np_dtype", "_canonical_name", "__weakref__") def __init__(self, np_dtype: npt.DTypeLike): - self._canonic_name = np_dtype + self._canonical_name = np_dtype self._np_dtype = np.dtype(np_dtype) def __repr__(self) -> str: diff --git a/array_api_strict/tests/test_creation_functions.py b/array_api_strict/tests/test_creation_functions.py index ee1263ae..3a89cf1d 100644 --- a/array_api_strict/tests/test_creation_functions.py +++ b/array_api_strict/tests/test_creation_functions.py @@ -247,7 +247,7 @@ def test_ones_etc(self, func, device): def test_ones_like_etc_correct(self, func): # float32 is preserved a = ones(2, dtype=float32) - device = Device('F32_device') + device = Device('no_float64') b = func(a, device=device) assert b.dtype == self.info.default_dtypes(device=device)["real floating"] assert b.device == device @@ -267,15 +267,15 @@ def test_ones_like_etc_incorrect(self, func): # incompatible dtype inferred from `a.dtype` with pytest.raises((TypeError, ValueError)): - func(a, device=Device('F32_device')) + func(a, device=Device('no_float64')) # `a.dtype` is compatible but the explicit dtype= argument is incompatible a = ones(2, dtype=float32) with pytest.raises((TypeError, ValueError)): - func(a, device=Device('F32_device'), dtype=float64) + func(a, device=Device('no_float64'), dtype=float64) def test_eye(self): - device = Device('F32_device') + device = Device('no_float64') a = eye(3, device=device) assert a.dtype == self.info.default_dtypes(device=device)["real floating"] assert a.device == device @@ -284,7 +284,7 @@ def test_eye(self): eye(3, device=device, dtype=float64) def test_linspace(self): - device = Device('F32_device') + device = Device('no_float64') a = linspace(1, 10, 11, device=device) assert a.dtype == self.info.default_dtypes(device=device)["real floating"] @@ -297,7 +297,7 @@ def test_linspace(self): linspace(1, 10, 11, device=device, dtype=float64) def test_arange(self): - device = Device('F32_device') + device = Device('no_float64') a = arange(0, 10, 1, device=device) assert a.dtype == self.info.default_dtypes(device=device)["integral"] @@ -314,7 +314,7 @@ def test_arange(self): arange(0.0, 10, 1, device=device, dtype=float64) def test_asarray(self): - device = Device('F32_device') + device = Device('no_float64') ### asarray(python_object) for x in (True, [False,]): diff --git a/array_api_strict/tests/test_device_support.py b/array_api_strict/tests/test_device_support.py index 84954db8..f2b0ba19 100644 --- a/array_api_strict/tests/test_device_support.py +++ b/array_api_strict/tests/test_device_support.py @@ -41,7 +41,7 @@ def test_fft_device_support_real(func_name): @pytest.mark.parametrize("func_name", ("fftfreq", "rfftfreq")) def test_fft_default_dtype(func_name): func = getattr(xp.fft, func_name) - device = xp.Device("F32_device") + device = xp.Device("no_float64") res = func(3, device=device) assert res.device == device assert res.dtype == xp.__array_namespace_info__().default_dtypes(device=device)["real floating"] @@ -53,23 +53,23 @@ def test_fft_default_dtype(func_name): class TestF32Device: @pytest.mark.parametrize("dtype_str", ["float64", "complex128"]) def test_f64_raises(self, dtype_str): - f32_device = xp.Device("F32_device") + f32_only_device = xp.Device("no_float64") dtype = getattr(xp, dtype_str) with pytest.raises(ValueError): - xp.arange(3, device=f32_device, dtype=dtype) + xp.arange(3, device=f32_only_device, dtype=dtype) def test_info_no_f64(self): - f32_device = xp.Device("F32_device") + f32_only_device = xp.Device("no_float64") info = xp.__array_namespace_info__() - all_dtypes = info.dtypes(device=f32_device) + all_dtypes = info.dtypes(device=f32_only_device) assert "float64" not in all_dtypes assert "complex128" not in all_dtypes def test_info_default_dtypes(self): - f32_device = xp.Device("F32_device") + f32_only_device = xp.Device("no_float64") info = xp.__array_namespace_info__() - defaults = info.default_dtypes(device=f32_device) + defaults = info.default_dtypes(device=f32_only_device) assert defaults["real floating"] == xp.float32 assert defaults["complex floating"] == xp.complex64 From a15c7bc6f930fa16cc909d6384d437579b6311fc Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Mon, 6 Jul 2026 17:03:30 +0200 Subject: [PATCH 09/10] raise if (device_type, device_id) is not found in the DLPack mappings --- array_api_strict/_devices.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index d23dcba3..41cb729d 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -62,9 +62,8 @@ def _normalize_dl_device(device_type: IntEnum | int, device_id: int) -> tuple[in def _device_from_dlpack_device( device_type: IntEnum | int, device_id: int ) -> Device: - return _DLPACK_DEVICE_TO_LOGICAL.get( - _normalize_dl_device(device_type, device_id), CPU_DEVICE - ) + # NB: if the (device_type, device_id) pair not known, raise + return _DLPACK_DEVICE_TO_LOGICAL[_normalize_dl_device(device_type, device_id)] def check_device(device: Device | None) -> None: From 572c4eb680039b677e3b6e79c68ec97d54cf7503 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Tue, 7 Jul 2026 12:52:00 +0200 Subject: [PATCH 10/10] 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. --- array_api_strict/_creation_functions.py | 31 +++++++++---------- array_api_strict/_devices.py | 9 ++++++ .../tests/test_creation_functions.py | 14 ++++++++- 3 files changed, 37 insertions(+), 17 deletions(-) diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index 52a0be47..b34af5ab 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -110,23 +110,22 @@ def asarray( # numpy default dtype may differ; if so, adjust the dtype if dtype is None and device is not None: res_dtype = DType(res.dtype) - if not device_supports_dtype(device, res_dtype): - # The dtype selected by Numpy might not be the default dtype - # on this device. If the dtype is not supported by the device we - # try to find one that is, aka the default dtype of this device. - from ._data_type_functions import isdtype - if isdtype(res_dtype, "bool"): - target_dtype = DType("bool") - elif isdtype(res_dtype, "integral"): - target_dtype = get_default_dtypes(device)["integral"] - elif isdtype(res_dtype, "real floating"): - target_dtype = get_default_dtypes(device)["real floating"] - elif isdtype(res_dtype, "complex floating"): - target_dtype = get_default_dtypes(device)["complex floating"] - else: - raise ValueError(f"{res_dtype = } not understood.") + # The dtype selected by Numpy might not be the default dtype + # on this device. We thus find the default dtype for the dtype "kind", and + # cast to the device-appropriate default. + from ._data_type_functions import isdtype + if isdtype(res_dtype, "bool"): + target_dtype = DType("bool") + elif isdtype(res_dtype, "integral"): + target_dtype = get_default_dtypes(device)["integral"] + elif isdtype(res_dtype, "real floating"): + target_dtype = get_default_dtypes(device)["real floating"] + elif isdtype(res_dtype, "complex floating"): + target_dtype = get_default_dtypes(device)["complex floating"] + else: + raise ValueError(f"{res_dtype = } not understood.") - res = res.astype(target_dtype._np_dtype) + res = res.astype(target_dtype._np_dtype) return Array._new(res, device=device) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index 41cb729d..c48db2a0 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -78,6 +78,15 @@ def check_device(device: Device | None) -> None: def get_default_dtypes(device: Device | None = None) -> dict[str, DType]: if device == NO_FLOAT64_DEVICE: + # mimic an MPS device which does not have float64 at all + return { + "real floating": float32, + "complex floating": complex64, + "integral": int64, + "indexing": int64, + } + elif device == Device('device2'): + # mimic a torch CPU device: support float64 but default to float32 return { "real floating": float32, "complex floating": complex64, diff --git a/array_api_strict/tests/test_creation_functions.py b/array_api_strict/tests/test_creation_functions.py index 3a89cf1d..70311206 100644 --- a/array_api_strict/tests/test_creation_functions.py +++ b/array_api_strict/tests/test_creation_functions.py @@ -22,7 +22,7 @@ zeros, zeros_like, ) -from .._dtypes import float32, float64, bool as xp_bool +from .._dtypes import float32, float64, complex64, bool as xp_bool from .._array_object import Array from .._devices import CPU_DEVICE, ALL_DEVICES, Device from .._info import __array_namespace_info__ @@ -356,6 +356,18 @@ def test_asarray(self): asarray(src, device=device) +def test_asarray_device_2(): + # device2 allows float64 but defaults to float32 + x = asarray([1.0], device=Device('device2')) + assert x.dtype == float32 + + x = asarray([1j], device=Device('device2')) + assert x.dtype == complex64 + + y = asarray([1.0], device=Device('device2'), dtype=float64) + assert y.dtype == float64 + + @pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12']) def from_dlpack_2023_12(api_version): if api_version != '2022.12':