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

Merged
ev-br merged 5 commits into from
Jul 27, 2026
Merged

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 ?

@betatim betatim Jul 27, 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.

The reason I thought that this test should "just work" is that both x and y are array-api-strict arrays. But spending two minutes staring at the newly failing test in your comment doesn't really help me understand why your proposed fix shouldn't work.

I think this means that this is more complicated than I think :D So fine to fix this in the future or never or realise that this isn't actually broken

@ev-br ev-br Jul 27, 2026 •
edited
Loading

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.

This one is confusing yes! And I had a head start to stare at this during EuroScipy. The bottom line AFAIU is this. Currently (in 2.5 and 2.6.1):

In [1]: import numpy as np

In [2]: import array_api_strict as xp

In [4]: a = np.ones(3)

In [5]: b = xp.from_dlpack(a)

In [6]: b.device
Out[6]: array_api_strict.Device('CPU_DEVICE')

In [7]: a.device
Out[7]: 'cpu'

The failed fix from #221 (comment) was checking hasattr(a, "device") which was a Yes, and then it attached a.device to b. But a.device is a string 'cpu' not an xp.Device instance.

Note that it currently works mainly by chance: xp.from_dlpack uses the default device unless explicitly provided as a kwarg:

In [8]: x = xp.ones(3, device=xp.Device("device1"))

In [9]: xp.from_dlpack(x).device
Out[9]: array_api_strict.Device('CPU_DEVICE')

So the pickle we're in for xp.from_dlpack(x) is this:

  • either we do if hasattr(x, "device"): device=x.device dance --- and then if x is a foreign array, we attach a foreign device to an xp-array;
  • or we do what 2.5 was doing, which is device=None means the array-api-strict default device---and then if x is an xp-array on a non-default device, we effectively transfer it to the default "CPU_DEVICE".

Your test was guarding for the latter, but the fix was breaking the former.

Now, the 2.5 and 2.6.1 behavior, strictly speaking, is questionable in view of the spec requirement (emphasis mine)

device on which to place the created array. If device is None and x supports DLPack, the output array must be on the same device as x. Default: None.

So which array-api-strict device is the same as NumPy's "cpu" device? And if "device1" is different, then how to tell it is different without somehow mapping the device values between different libraries?

@betatim betatim Jul 28, 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.

Naively (as an unsuspecting user) the behaviour that the standard describes makes sense. It is what I'd expect to happen.

That the device changes in:

In [8]: x = xp.ones(3, device=xp.Device("device1"))

In [9]: xp.from_dlpack(x).device
Out[9]: array_api_strict.Device('CPU_DEVICE')

seems super weird as a user because we don't change array library in the process.

I think attaching the device from a different array library to the result of from_dlpack is the wrong thing to do. The standard says that how a device is specified is an implementation detail of the library (some use a string, some a class, etc). So attaching a pytorch device to a cupy array won't do anything useful, as none of the cupy code will understand it.

I think the question of "how do you translate a device?" is the key. Can we use the device type and ID of the source array to find the corresponding device in array-api-strict?

In [3]: x = torch.asarray([1.,2,.3])

In [4]: x.__dlpack_device__()
Out[4]: (<DLDeviceType.kDLCPU: 1>, 0)

In [5]: x = torch.asarray([1.,2,.3], device="mps")

In [6]: x.__dlpack_device__()
Out[6]: (<DLDeviceType.kDLMetal: 8>, 0)

So I'd adjust the device handling in

array-api-strict/array_api_strict/_creation_functions.py

Line 234 in f02a5da

def from_dlpack(
to call __dlpack_device__() and see if there is a matching device in array-api-strict if the argument device=None. My intuition would be that __dlpack_device__ returns (1, 0) for Device("CPU_DEVICE") and (1, 1) for Device("device1"). Or something like that. As we learnt we can't change the type because that screws up things, but I think we can use the device ID to tell these otherwise similar devices apart.

@ev-br ev-br Jul 28, 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.

Agreed that the current behaviour is weird! And agreed about the key question, too. Any chance you've cycles to spare to prototype a solution? You seem to have a setup for downstream testing (in scikit-learn or where was that the previous attempt wrought havoc in?).

"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