Skip to content

BUG: Revert dlpack device #221

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ev-br wants to merge 5 commits into
base: main
Choose a base branch
from
Open

BUG: Revert dlpack device #221

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
744a5d4
Revert "raise if (device_type, device_id) is not found in the DLPack …
ev-br Jul 10, 2026
f92f9ba
Revert "add dlpack_device dunder to the F32_device"
ev-br Jul 10, 2026
43edaaa
Revert "manually sync dlpack enum updates to _devices.py"
ev-br Jul 10, 2026
1178e46
Manually fixup reverting dlpack_device additions
ev-br Jul 10, 2026
21b4404
MAINT: remove a stale comment
ev-br Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 9 additions & 30 deletions array_api_strict/_array_object.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand All @@ -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]
"""
Expand Down
4 changes: 0 additions & 4 deletions array_api_strict/_creation_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,11 +250,7 @@ def from_dlpack(
_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)
Expand Down
30 changes: 0 additions & 30 deletions array_api_strict/_devices.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
from typing import Final
from enum import IntEnum

from ._dtypes import (
DType, float32, float64, complex64, complex128, int64,
Expand Down Expand Up @@ -36,35 +35,6 @@ def __hash__(self) -> int:

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)]


def check_device(device: Device | None) -> None:
if device is not None and not isinstance(device, Device):
Expand Down
35 changes: 1 addition & 34 deletions array_api_strict/tests/test_array_object.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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__`)"""
Expand Down
9 changes: 0 additions & 9 deletions array_api_strict/tests/test_creation_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -402,12 +402,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(

@betatim betatim Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could keep this test no? It seems like it checks something that is somewhat related to the DLPack device ID but also not related.

I'd keep it (unless I'm missing something)

@ev-br ev-br Jul 22, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An issue is that this test just did not work in 2.5? So if we're just reverting the change made in 2.6, the test is broken; if we also fix from_dlpack, that's not just a revert...

In [1]: import array_api_strict as xp

In [2]: xp.__version__
Out[2]: '2.5'

In [3]: a = xp.arange(3.0, device=xp.Device('device1'))

In [4]: xp.from_dlpack(a).device
Out[4]: array_api_strict.Device('CPU_DEVICE')

An attempt at a fix:

$ git diff
diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py
index 685044d..5e40771 100644
--- a/array_api_strict/_creation_functions.py
+++ b/array_api_strict/_creation_functions.py
@@ -250,6 +250,8 @@ def from_dlpack(
         _check_device(device)
     else:
         device = None
+        if hasattr(x, "device"):
+            device = x.device
 
     if copy in [_undef, None]:
         # numpy 1.26 does not have the copy= arg

Looks simple and correct, right?

array_api_strict/tests/test_creation_functions.py::test_from_dlpack_default_device FAILED                           [ 96%]
>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> traceback >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

    def test_from_dlpack_default_device():
        x = asarray([1, 2, 3])
        y = from_dlpack(x)
        z = from_dlpack(np.asarray([1, 2, 3]))
>       assert x.device == y.device == z.device == CPU_DEVICE
E       AssertionError: assert array_api_strict.Device('CPU_DEVICE') == 'cpu'
E        +  where array_api_strict.Device('CPU_DEVICE') = Array([1, 2, 3], dtype=array_api_strict.int64).device
E        +  and   'cpu' = Array([1, 2, 3], dtype=array_api_strict.int64, device=cpu).device

So one cannot just reuse the device of x, one has to map the device of x onto the device of array-api-strict! And the device naming is not standardized, so pretty much the only way to map them is via DLPACK enum value, but that's what we reverting here.

TBH I don't see a way out, do you @betatim ?

"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
Loading