BUG: Revert dlpack device

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

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)

array_api_strict/_devices.py

@@ -1,5 +1,4 @@
 from typing import Final
-from enum import IntEnum
 
 from ._dtypes import (
     DType, float32, float64, complex64, complex128, int64,
@@ -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):

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__`)"""

array_api_strict/tests/test_creation_functions.py

@@ -395,12 +395,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