diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py index 8420aee..b0689e1 100644 --- a/array_api_strict/_array_object.py +++ b/array_api_strict/_array_object.py @@ -40,10 +40,7 @@ _real_to_complex_map, _result_type, ) -from ._devices import ( - CPU_DEVICE, Device, device_supports_dtype, _normalize_dl_device, _DLPACK_DEVICE_FOR, - DLDeviceType -) +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 @@ -614,31 +611,12 @@ def __dlpack__( raise NotImplementedError("The copy argument to __dlpack__ is not yet implemented") return self._array.__dlpack__(stream=stream) - - kwargs: dict[str, Any] = {'stream': stream} - self_dl_device = _normalize_dl_device(*_DLPACK_DEVICE_FOR[self._device]) - cpu_dl_device = _normalize_dl_device(DLDeviceType.kDLCPU, 0) - numpy_dl_device: tuple[IntEnum, int] | None = None - - if dl_device not in [_undef, None]: - requested = _normalize_dl_device(dl_device[0], dl_device[1]) - if requested == self_dl_device: - pass - elif requested == cpu_dl_device and self_dl_device != cpu_dl_device: - if copy is False: - raise BufferError( - "Cannot export array to CPU without copying when copy=False" - ) - if copy is _undef: - copy = True - numpy_dl_device = (DLDeviceType.kDLCPU, 0) - else: - raise BufferError("unsupported device requested") - - if max_version is not _undef: - kwargs['max_version'] = max_version - if numpy_dl_device is not None: - kwargs['dl_device'] = numpy_dl_device + else: + kwargs = {'stream': stream} + if max_version is not _undef: + kwargs['max_version'] = max_version + if dl_device is not _undef: + kwargs['dl_device'] = dl_device if copy is not _undef: kwargs['copy'] = copy return self._array.__dlpack__(**kwargs) @@ -647,7 +625,8 @@ def __dlpack_device__(self) -> tuple[IntEnum, int]: """ Performs the operation __dlpack_device__. """ - return _DLPACK_DEVICE_FOR[self._device] + # Note: device support is required for this + return self._array.__dlpack_device__() def __eq__(self, other: Array | complex, /) -> Array: # type: ignore[override] """ diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index b34af5a..0f8fc51 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -242,15 +242,12 @@ def from_dlpack( if copy is not _undef: raise ValueError("The copy argument to from_dlpack requires at least version 2023.12 of the array API") + # Going to wait for upstream numpy support if device is not _undef: _check_device(device) else: device = None - if hasattr(x, "__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) if copy in [_undef, None]: # numpy 1.26 does not have the copy= arg return Array._new(np.from_dlpack(x), device=device) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index c48db2a..1883199 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -1,14 +1,13 @@ from typing import Final -from enum import IntEnum from ._dtypes import ( - DType, float32, float64, complex64, complex128, int64, + DType, float32, float64, complex64, complex128, int64, uint64, int32, _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", "no_float64") +_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "no_float64", "no_x64") class Device: _device: Final[str] @@ -33,37 +32,11 @@ def __hash__(self) -> int: CPU_DEVICE = Device() NO_FLOAT64_DEVICE = Device("no_float64") +NO_X64_DEVICE = Device("no_x64") -ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), NO_FLOAT64_DEVICE) - -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), - NO_FLOAT64_DEVICE: (DLDeviceType.kDLMETAL, 0), -} - -_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: - # NB: if the (device_type, device_id) pair not known, raise - return _DLPACK_DEVICE_TO_LOGICAL[_normalize_dl_device(device_type, device_id)] +ALL_DEVICES = ( + CPU_DEVICE, Device("device1"), Device("device2"), NO_FLOAT64_DEVICE, NO_X64_DEVICE +) def check_device(device: Device | None) -> None: @@ -85,6 +58,14 @@ def get_default_dtypes(device: Device | None = None) -> dict[str, DType]: "integral": int64, "indexing": int64, } + elif device == NO_X64_DEVICE: + # mimic JAX default: no float64, no int64 + return { + "real floating": float32, + "complex floating": complex64, + "integral": int32, + "indexing": int32, + } elif device == Device('device2'): # mimic a torch CPU device: support float64 but default to float32 return { @@ -107,6 +88,9 @@ def device_supports_dtype(device: Device | None, dtype: DType |None) -> bool: # Device("no_float64") supports all dtypes except float64 and complex128 if device == NO_FLOAT64_DEVICE: return dtype not in (float64, complex128) + if device == NO_X64_DEVICE: + # "no_x64" only supports 32-bit ints and floats + return dtype not in (float64, complex128, int64, uint64) # All other devices support all dtypes return True diff --git a/array_api_strict/tests/test_array_object.py b/array_api_strict/tests/test_array_object.py index a75d66f..445ee80 100644 --- a/array_api_strict/tests/test_array_object.py +++ b/array_api_strict/tests/test_array_object.py @@ -10,7 +10,7 @@ from .. import ones, arange, reshape, asarray, result_type, all, equal, stack from .._array_object import Array -from .._devices import CPU_DEVICE, Device, DLDeviceType +from .._devices import CPU_DEVICE, Device from .._dtypes import ( _all_dtypes, _boolean_dtypes, @@ -758,39 +758,6 @@ def test_dlpack_2023_12(api_version): a.__dlpack__(copy=True) a.__dlpack__(copy=None) -@pytest.mark.parametrize( - ("device", "expected"), - [ - (CPU_DEVICE, (DLDeviceType.kDLCPU, 0)), - (Device("device1"), (DLDeviceType.kDLCUDA, 0)), - (Device("device2"), (DLDeviceType.kDLCUDA, 1)), - ], -) -def test_dlpack_device_numbers(device, expected): - a = asarray([1, 2, 3], device=device) - assert a.__dlpack_device__() == expected - - -@pytest.mark.parametrize("device", [CPU_DEVICE, Device("device1"), Device("device2")]) -def test_dlpack_export_with_matching_dl_device(device): - if np.lib.NumpyVersion(np.__version__) < "2.1.0": - pytest.skip("dl_device argument requires NumPy >= 2.1.0") - set_array_api_strict_flags(api_version="2023.12") - - a = asarray([1, 2, 3], device=device) - a.__dlpack__(dl_device=a.__dlpack_device__()) - - -@pytest.mark.parametrize("device", [Device("device1"), Device("device2")]) -def test_dlpack_cross_device_export_buffer_error(device): - if np.lib.NumpyVersion(np.__version__) < "2.1.0": - pytest.skip("dl_device argument requires NumPy >= 2.1.0") - set_array_api_strict_flags(api_version="2023.12") - - a = asarray([1, 2, 3], device=device) - with pytest.raises(BufferError): - a.__dlpack__(dl_device=(DLDeviceType.kDLCPU, 0), copy=False) - def test_pickle(): """Check that arrays are pickleable (despite raising on `__new__`)""" diff --git a/array_api_strict/tests/test_creation_functions.py b/array_api_strict/tests/test_creation_functions.py index 7031120..8a480bb 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, complex64, bool as xp_bool +from .._dtypes import float32, float64, complex64, int32, int64, bool as xp_bool from .._array_object import Array from .._devices import CPU_DEVICE, ALL_DEVICES, Device from .._info import __array_namespace_info__ @@ -368,6 +368,17 @@ def test_asarray_device_2(): assert y.dtype == float64 +def test_asarray_no_x64_device(): + x = asarray(3, device=Device("no_x64")) + assert x.dtype == int32 + + with pytest.raises(ValueError): + asarray(3, device=Device("no_x64"), dtype=int64) + + y = zeros_like(ones(3, dtype=int32), device=Device("no_x64")) + assert y.dtype == int32 + + @pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12']) def from_dlpack_2023_12(api_version): if api_version != '2022.12': @@ -395,12 +406,3 @@ def test_from_dlpack_default_device(): z = from_dlpack(np.asarray([1, 2, 3])) assert x.device == y.device == z.device == CPU_DEVICE - -@pytest.mark.parametrize( - "device", - [Device("device1"), Device("device2")], -) -def test_from_dlpack_preserves_device(device): - x = asarray([1, 2, 3], device=device) - y = from_dlpack(x) - assert y.device == device