From 59556d8723a1e0f5a4dfadd4ba01ee4ff5398ee4 Mon Sep 17 00:00:00 2001 From: Tim Head Date: Wed, 29 Jul 2026 13:58:15 +0200 Subject: [PATCH 1/3] Improve DLPack support for non CPU devices This improves the `from_dlpack` behaviour for round-tripping array-api-strict arrays and raises an error for arrays that are not on the CPU device. --- .github/workflows/tests.yml | 12 +++- array-api-tests-xfails.txt | 4 -- array_api_strict/_array_object.py | 40 ++++++++--- array_api_strict/_creation_functions.py | 11 ++- array_api_strict/_devices.py | 61 ++++++++++++++++ array_api_strict/tests/test_array_object.py | 69 ++++++++++++++++++- .../tests/test_creation_functions.py | 28 +++++++- array_api_strict/tests/test_dlpack_interop.py | 44 ++++++++++++ 8 files changed, 252 insertions(+), 17 deletions(-) create mode 100644 array_api_strict/tests/test_dlpack_interop.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 03cdd919..ab2dac27 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -11,11 +11,16 @@ jobs: - python-version: '3.10' numpy-version: 'latest' - python-version: '3.10' - numpy-version: 'dev' + numpy-version: 'dev' - python-version: '3.13' numpy-version: '1.26' - python-version: '3.14' numpy-version: '1.26' + include: + # exercise the DLPack interop tests, which need another array library + - python-version: '3.13' + numpy-version: 'latest' + torch: true fail-fast: false steps: - uses: actions/checkout@v7.0.1 @@ -34,6 +39,11 @@ jobs: fi python -m pip install pytest hypothesis python -c'import numpy as np; print(f"{np.__version__ = }")' + - name: Install PyTorch + if: matrix.torch + run: | + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -c'import torch; print(f"{torch.__version__ = }")' - name: Run Tests run: | pytest diff --git a/array-api-tests-xfails.txt b/array-api-tests-xfails.txt index b7042ffe..21a43426 100644 --- a/array-api-tests-xfails.txt +++ b/array-api-tests-xfails.txt @@ -46,9 +46,5 @@ array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is +0 and imag array_api_tests/test_special_cases.py::test_unary[sqrt(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +0 + infinity j] -# remove this when array-api-strict 2.6 is released: this is only for array-api-tests, -# which runs self-tests against a released array-api-strict -array_api_tests/test_dlpack.py::test_from_dlpack - # NumPy 1.26 : NotImplementedError: The copy argument to __dlpack__ is not yet implemented array_api_tests/test_dlpack.py::test_dunder_dlpack diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py index 25e6595a..e8b4e27d 100644 --- a/array_api_strict/_array_object.py +++ b/array_api_strict/_array_object.py @@ -40,7 +40,13 @@ _real_to_complex_map, _result_type, ) -from ._devices import CPU_DEVICE, Device, device_supports_dtype +from ._devices import ( + CPU_DEVICE, + Device, + _DLPACK_DEVICE_FOR, + _normalize_dl_device, + device_supports_dtype, +) from ._flags import get_array_api_strict_flags, set_array_api_strict_flags from ._typing import PyCapsule @@ -602,6 +608,15 @@ def __dlpack__( if copy is not _undef: raise ValueError("The copy argument to __dlpack__ requires at least version 2023.12 of the array API") + if self._device != CPU_DEVICE: + # Consistent with __buffer__ and __array__: data cannot leave a + # non-CPU device implicitly. Note that this also rules out consumers + # asking for the CPU device via dl_device: the array has to be moved + # to the CPU device first, with to_device or asarray. + raise BufferError( + f"Can't export array on the '{self._device}' device via DLPack." + ) + if np.lib.NumpyVersion(np.__version__) < '2.1.0': if max_version not in [_undef, None]: raise NotImplementedError("The max_version argument to __dlpack__ is not yet implemented") @@ -611,12 +626,20 @@ def __dlpack__( raise NotImplementedError("The copy argument to __dlpack__ is not yet implemented") return self._array.__dlpack__(stream=stream) - 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 + + kwargs: dict[str, Any] = {'stream': stream} + if max_version is not _undef: + kwargs['max_version'] = max_version + if dl_device is not _undef: + if dl_device is not None: + requested = _normalize_dl_device(*dl_device) + if requested != _normalize_dl_device(*_DLPACK_DEVICE_FOR[self._device]): + raise BufferError("unsupported device requested") + # The request is for the device the array is already on, which a + # plain export satisfies. NumPy is not told about it, as it only + # knows about its own spelling of the CPU device. + dl_device = None + kwargs['dl_device'] = dl_device if copy is not _undef: kwargs['copy'] = copy return self._array.__dlpack__(**kwargs) @@ -625,8 +648,7 @@ def __dlpack_device__(self) -> tuple[IntEnum, int]: """ Performs the operation __dlpack_device__. """ - # Note: device support is required for this - return self._array.__dlpack_device__() + return _DLPACK_DEVICE_FOR[self._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 685044df..eb70ca4d 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -7,7 +7,7 @@ from ._dtypes import DType, _all_dtypes, _np_dtype, bool as xp_bool from ._devices import ( - Device, device_supports_dtype, get_default_dtypes, + Device, device_supports_dtype, get_default_dtypes, _device_from_dlpack_device, check_device as _check_device ) from ._flags import get_array_api_strict_flags @@ -250,6 +250,15 @@ def from_dlpack( _check_device(device) else: device = None + if hasattr(x, "__dlpack_device__"): + dl_type, dl_id = x.__dlpack_device__() + device = _device_from_dlpack_device(dl_type, dl_id) + + if isinstance(x, Array): + # Arrays of this library are unwrapped instead of going through DLPack: + # the buffer is in host memory whatever the logical device is, and the + # DLPack export refuses arrays which are not on the CPU device. + x = x._array if copy in [_undef, None]: # numpy 1.26 does not have the copy= arg diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index 18831991..d2ee667f 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -1,3 +1,4 @@ +from enum import IntEnum from typing import Final from ._dtypes import ( @@ -39,6 +40,66 @@ def __hash__(self) -> int: ) +class DLDeviceType(IntEnum): + """The DLPack device types, as defined by the DLPack ABI.""" + kDLCPU = 1 + kDLCUDA = 2 + kDLCUDAHost = 3 + kDLOpenCL = 4 + kDLVulkan = 7 + kDLMetal = 8 + kDLVPI = 9 + kDLROCM = 10 + kDLROCMHost = 11 + kDLExtDev = 12 + kDLCUDAManaged = 13 + kDLOneAPI = 14 + kDLWebGPU = 15 + kDLHexagon = 16 + kDLMAIA = 17 + + +# All the devices of array_api_strict are fictitious and their data lives in host +# memory, so they all report the CPU device type and are told apart by their device +# id. Reporting anything else makes consumers such as pytorch dispatch to the +# machinery of a device they cannot actually reach, see +# https://github.com/data-apis/array-api-strict/issues/219 +_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = { + CPU_DEVICE: (DLDeviceType.kDLCPU, 0), + Device("device1"): (DLDeviceType.kDLCPU, 1), + Device("device2"): (DLDeviceType.kDLCPU, 2), + NO_FLOAT64_DEVICE: (DLDeviceType.kDLCPU, 3), + NO_X64_DEVICE: (DLDeviceType.kDLCPU, 4), +} + +_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]: + # `device_type` may be a member of another library's DLPack enum + return (int(device_type), device_id) + + +def _device_from_dlpack_device( + device_type: IntEnum | int, device_id: int +) -> Device: + key = _normalize_dl_device(device_type, device_id) + try: + return _DLPACK_DEVICE_TO_LOGICAL[key] + except KeyError: + try: + type_name = DLDeviceType(key[0]).name + except ValueError: + type_name = str(key[0]) + raise BufferError( + f"No array_api_strict device matches the DLPack device " + f"({type_name}, {device_id})." + ) from None + + def check_device(device: Device | None) -> None: if device is not None and not isinstance(device, Device): raise ValueError(f"Unsupported device {device!r}") diff --git a/array_api_strict/tests/test_array_object.py b/array_api_strict/tests/test_array_object.py index 445ee80b..655b17d9 100644 --- a/array_api_strict/tests/test_array_object.py +++ b/array_api_strict/tests/test_array_object.py @@ -10,7 +10,9 @@ from .. import ones, arange, reshape, asarray, result_type, all, equal, stack from .._array_object import Array -from .._devices import CPU_DEVICE, Device +from .._devices import ( + ALL_DEVICES, CPU_DEVICE, Device, DLDeviceType, _DLPACK_DEVICE_FOR +) from .._dtypes import ( _all_dtypes, _boolean_dtypes, @@ -759,6 +761,71 @@ def test_dlpack_2023_12(api_version): a.__dlpack__(copy=None) +@pytest.mark.parametrize( + ("device", "expected"), + [ + (CPU_DEVICE, (DLDeviceType.kDLCPU, 0)), + (Device("device1"), (DLDeviceType.kDLCPU, 1)), + (Device("device2"), (DLDeviceType.kDLCPU, 2)), + (Device("no_float64"), (DLDeviceType.kDLCPU, 3)), + (Device("no_x64"), (DLDeviceType.kDLCPU, 4)), + ], +) +def test_dlpack_device_numbers(device, expected): + a = asarray([1, 2, 3], device=device) + assert a.__dlpack_device__() == expected + + +def test_dlpack_device_map_is_complete(): + assert set(_DLPACK_DEVICE_FOR) == set(ALL_DEVICES) + # devices which share a DLPack device cannot be told apart by from_dlpack + assert len(set(_DLPACK_DEVICE_FOR.values())) == len(ALL_DEVICES) + + +@pytest.mark.parametrize( + "device", [device for device in ALL_DEVICES if device != CPU_DEVICE] +) +def test_dlpack_export_from_non_cpu_device(device): + a = asarray([1, 2, 3], device=device) + + with pytest.raises(BufferError): + a.__dlpack__() + with pytest.raises(BufferError): + np.from_dlpack(a) + + if np.lib.NumpyVersion(np.__version__) < "2.1.0": + return + + # asking for a device explicitly does not help: the array has to be moved + # to the CPU device first. Even asking for the device the array is already + # on is refused, since the capsule can only ever be tagged with the CPU + # device and the consumer would end up with the data on the wrong device. + # array-api-strict's special devices all use the same DLPack device number, + # but are meant to represent devices like a GPU or other accelerator. + with pytest.raises(BufferError): + a.__dlpack__(dl_device=a.__dlpack_device__()) + with pytest.raises(BufferError): + a.__dlpack__(dl_device=(DLDeviceType.kDLCPU, 0)) + with pytest.raises(BufferError): + np.from_dlpack(a, device="cpu") + + np.from_dlpack(a.to_device(CPU_DEVICE)) + + +def test_dlpack_export_from_cpu_device(): + a = asarray([1, 2, 3]) + + a.__dlpack__() + np.from_dlpack(a) + + if np.lib.NumpyVersion(np.__version__) < "2.1.0": + return + + a.__dlpack__(dl_device=a.__dlpack_device__()) + with pytest.raises(BufferError): + a.__dlpack__(dl_device=(DLDeviceType.kDLCUDA, 0)) + + def test_pickle(): """Check that arrays are pickleable (despite raising on `__new__`)""" a = ones(2) diff --git a/array_api_strict/tests/test_creation_functions.py b/array_api_strict/tests/test_creation_functions.py index a28be649..4cbce07a 100644 --- a/array_api_strict/tests/test_creation_functions.py +++ b/array_api_strict/tests/test_creation_functions.py @@ -24,7 +24,7 @@ ) 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 .._devices import CPU_DEVICE, ALL_DEVICES, Device, DLDeviceType from .._info import __array_namespace_info__ from .._flags import set_array_api_strict_flags @@ -413,3 +413,29 @@ 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", ALL_DEVICES) +def test_from_dlpack_preserves_device(device): + x = asarray([1, 2, 3], device=device) + y = from_dlpack(x) + assert y.device == device + + +def test_from_dlpack_unknown_device(): + class ForeignArray: + """An array on a device which array_api_strict knows nothing about.""" + def __init__(self): + self._array = np.asarray([1, 2, 3]) + + def __dlpack_device__(self): + return (DLDeviceType.kDLCUDA, 0) + + def __dlpack__(self, **kwargs): + return self._array.__dlpack__(**kwargs) + + with pytest.raises(BufferError): + from_dlpack(ForeignArray()) + + # an explicit device is a request to transfer, so nothing has to be inferred + assert from_dlpack(ForeignArray(), device=CPU_DEVICE).device == CPU_DEVICE + diff --git a/array_api_strict/tests/test_dlpack_interop.py b/array_api_strict/tests/test_dlpack_interop.py new file mode 100644 index 00000000..4ac15c3e --- /dev/null +++ b/array_api_strict/tests/test_dlpack_interop.py @@ -0,0 +1,44 @@ +"""Check how other libraries see arrays of array_api_strict via DLPack. + +Tests that libraries like NumPy and PyTorch can import arrays from +array_api_strict via DLPack. + +Only run if torch is available. +""" +import numpy as np +import pytest + +import array_api_strict as xp +from .._devices import ALL_DEVICES, CPU_DEVICE + +torch = pytest.importorskip("torch") + + +def test_export_from_cpu_device(): + x = xp.asarray([1, 2, 3], device=CPU_DEVICE) + + assert np.from_dlpack(x).tolist() == [1, 2, 3] + assert torch.from_dlpack(x).tolist() == [1, 2, 3] + + +@pytest.mark.parametrize( + "device", [device for device in ALL_DEVICES if device != CPU_DEVICE] +) +def test_export_from_other_devices(device): + x = xp.asarray([1, 2, 3], device=device) + + with pytest.raises(BufferError): + np.from_dlpack(x) + with pytest.raises(BufferError): + torch.from_dlpack(x) + + assert torch.from_dlpack(x.to_device(CPU_DEVICE)).tolist() == [1, 2, 3] + + +@pytest.mark.parametrize("device", ALL_DEVICES) +def test_import_from_torch(device): + # int32 is the widest integer every device supports + x = xp.from_dlpack(torch.asarray([1, 2, 3], dtype=torch.int32), device=device) + + assert x.device == device + assert xp.all(x == xp.asarray([1, 2, 3], dtype=xp.int32, device=device)) From fa2780eb6d901d240ea3fc7d8d4248aecd77e0d8 Mon Sep 17 00:00:00 2001 From: Tim Head Date: Wed, 29 Jul 2026 15:19:03 +0200 Subject: [PATCH 2/3] Remove using different device IDs for logical devices The CPU device can not really have multiple devices, and we don't need to be able to tell devices apart. So this removes the "not spec compliant" code. --- array_api_strict/_creation_functions.py | 6 ++++- array_api_strict/_devices.py | 18 ++++++------- array_api_strict/tests/test_array_object.py | 28 +++++++-------------- 3 files changed, 22 insertions(+), 30 deletions(-) diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index eb70ca4d..54e91e53 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -250,7 +250,11 @@ def from_dlpack( _check_device(device) else: device = None - if hasattr(x, "__dlpack_device__"): + if isinstance(x, Array): + # All the devices of this library share a DLPack device, so the + # logical device is read off the array itself. + device = x.device + elif hasattr(x, "__dlpack_device__"): dl_type, dl_id = x.__dlpack_device__() device = _device_from_dlpack_device(dl_type, dl_id) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index d2ee667f..79d8bdc8 100644 --- a/array_api_strict/_devices.py +++ b/array_api_strict/_devices.py @@ -60,21 +60,19 @@ class DLDeviceType(IntEnum): # All the devices of array_api_strict are fictitious and their data lives in host -# memory, so they all report the CPU device type and are told apart by their device -# id. Reporting anything else makes consumers such as pytorch dispatch to the -# machinery of a device they cannot actually reach, see +# memory, so they all report the CPU device, which is device number zero. Reporting +# anything else makes consumers such as pytorch dispatch to the machinery of a +# device they cannot actually reach, see # https://github.com/data-apis/array-api-strict/issues/219 +# The logical devices cannot be told apart through DLPack, which is fine: only the +# CPU device can be exported at all, and from_dlpack recovers the logical device of +# an array from this library without going through DLPack. _DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = { - CPU_DEVICE: (DLDeviceType.kDLCPU, 0), - Device("device1"): (DLDeviceType.kDLCPU, 1), - Device("device2"): (DLDeviceType.kDLCPU, 2), - NO_FLOAT64_DEVICE: (DLDeviceType.kDLCPU, 3), - NO_X64_DEVICE: (DLDeviceType.kDLCPU, 4), + device: (DLDeviceType.kDLCPU, 0) for device in ALL_DEVICES } _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() + (int(DLDeviceType.kDLCPU), 0): CPU_DEVICE, } diff --git a/array_api_strict/tests/test_array_object.py b/array_api_strict/tests/test_array_object.py index 655b17d9..8a048c4d 100644 --- a/array_api_strict/tests/test_array_object.py +++ b/array_api_strict/tests/test_array_object.py @@ -761,25 +761,16 @@ def test_dlpack_2023_12(api_version): a.__dlpack__(copy=None) -@pytest.mark.parametrize( - ("device", "expected"), - [ - (CPU_DEVICE, (DLDeviceType.kDLCPU, 0)), - (Device("device1"), (DLDeviceType.kDLCPU, 1)), - (Device("device2"), (DLDeviceType.kDLCPU, 2)), - (Device("no_float64"), (DLDeviceType.kDLCPU, 3)), - (Device("no_x64"), (DLDeviceType.kDLCPU, 4)), - ], -) -def test_dlpack_device_numbers(device, expected): +@pytest.mark.parametrize("device", ALL_DEVICES) +def test_dlpack_device_numbers(device): a = asarray([1, 2, 3], device=device) - assert a.__dlpack_device__() == expected + # the data of every logical device lives in host memory, so they all report + # the CPU device, which is device number zero + assert a.__dlpack_device__() == (DLDeviceType.kDLCPU, 0) def test_dlpack_device_map_is_complete(): assert set(_DLPACK_DEVICE_FOR) == set(ALL_DEVICES) - # devices which share a DLPack device cannot be told apart by from_dlpack - assert len(set(_DLPACK_DEVICE_FOR.values())) == len(ALL_DEVICES) @pytest.mark.parametrize( @@ -797,11 +788,10 @@ def test_dlpack_export_from_non_cpu_device(device): return # asking for a device explicitly does not help: the array has to be moved - # to the CPU device first. Even asking for the device the array is already - # on is refused, since the capsule can only ever be tagged with the CPU - # device and the consumer would end up with the data on the wrong device. - # array-api-strict's special devices all use the same DLPack device number, - # but are meant to represent devices like a GPU or other accelerator. + # to the CPU device first. Even asking for the CPU device, which is what + # __dlpack_device__ reports, is refused: these devices are meant to + # represent a GPU or other accelerator, so the consumer would end up with + # the data on a device the array is not logically on. with pytest.raises(BufferError): a.__dlpack__(dl_device=a.__dlpack_device__()) with pytest.raises(BufferError): From b9a87e0952eb20a5c486bfd31aa3e8875a9e3b18 Mon Sep 17 00:00:00 2001 From: Tim Head Date: Wed, 29 Jul 2026 16:33:45 +0200 Subject: [PATCH 3/3] Fix copy kwarg handling for old numpy versions --- array_api_strict/_creation_functions.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index 54e91e53..ca94e4d8 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -264,9 +264,14 @@ def from_dlpack( # DLPack export refuses arrays which are not on the CPU device. x = x._array - if copy in [_undef, None]: - # numpy 1.26 does not have the copy= arg - return Array._new(np.from_dlpack(x), device=device) + if copy is _undef: + copy = None + + if np.lib.NumpyVersion(np.__version__) < '2.1.0': + # numpy 1.26 does not have the copy= arg, and its from_dlpack never + # copies: copy=False needs nothing extra, copy=True is done here. + out = np.from_dlpack(x) + return Array._new(np.copy(out) if copy else out, device=device) return Array._new(np.from_dlpack(x, copy=copy), device=device)