Skip to content

Commit dd2ae23

Browse files
committed
Manually fixup reverting dlpack_device additions
1 parent 1efbf5e commit dd2ae23

3 files changed

Lines changed: 10 additions & 69 deletions

File tree

array_api_strict/_array_object.py

Lines changed: 9 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -40,10 +40,7 @@
4040
_real_to_complex_map,
4141
_result_type,
4242
)
43-
from ._devices import (
44-
CPU_DEVICE, Device, device_supports_dtype, _normalize_dl_device, _DLPACK_DEVICE_FOR,
45-
DLDeviceType
46-
)
43+
from ._devices import CPU_DEVICE, Device, device_supports_dtype
4744
from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
4845
from ._typing import PyCapsule
4946

@@ -614,31 +611,12 @@ def __dlpack__(
614611
raise NotImplementedError("The copy argument to __dlpack__ is not yet implemented")
615612

616613
return self._array.__dlpack__(stream=stream)
617-
618-
kwargs: dict[str, Any] = {'stream': stream}
619-
self_dl_device = _normalize_dl_device(*_DLPACK_DEVICE_FOR[self._device])
620-
cpu_dl_device = _normalize_dl_device(DLDeviceType.kDLCPU, 0)
621-
numpy_dl_device: tuple[IntEnum, int] | None = None
622-
623-
if dl_device not in [_undef, None]:
624-
requested = _normalize_dl_device(dl_device[0], dl_device[1])
625-
if requested == self_dl_device:
626-
pass
627-
elif requested == cpu_dl_device and self_dl_device != cpu_dl_device:
628-
if copy is False:
629-
raise BufferError(
630-
"Cannot export array to CPU without copying when copy=False"
631-
)
632-
if copy is _undef:
633-
copy = True
634-
numpy_dl_device = (DLDeviceType.kDLCPU, 0)
635-
else:
636-
raise BufferError("unsupported device requested")
637-
638-
if max_version is not _undef:
639-
kwargs['max_version'] = max_version
640-
if numpy_dl_device is not None:
641-
kwargs['dl_device'] = numpy_dl_device
614+
else:
615+
kwargs = {'stream': stream}
616+
if max_version is not _undef:
617+
kwargs['max_version'] = max_version
618+
if dl_device is not _undef:
619+
kwargs['dl_device'] = dl_device
642620
if copy is not _undef:
643621
kwargs['copy'] = copy
644622
return self._array.__dlpack__(**kwargs)
@@ -647,7 +625,8 @@ def __dlpack_device__(self) -> tuple[IntEnum, int]:
647625
"""
648626
Performs the operation __dlpack_device__.
649627
"""
650-
return _DLPACK_DEVICE_FOR[self._device]
628+
# Note: device support is required for this
629+
return self._array.__dlpack_device__()
651630

652631
def __eq__(self, other: Array | complex, /) -> Array: # type: ignore[override]
653632
"""

array_api_strict/_creation_functions.py

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -246,11 +246,6 @@ def from_dlpack(
246246
_check_device(device)
247247
else:
248248
device = None
249-
if hasattr(x, "__dlpack_device__"):
250-
from ._devices import _device_from_dlpack_device
251-
252-
dl_type, dl_id = x.__dlpack_device__()
253-
device = _device_from_dlpack_device(dl_type, dl_id)
254249
if copy in [_undef, None]:
255250
# numpy 1.26 does not have the copy= arg
256251
return Array._new(np.from_dlpack(x), device=device)

array_api_strict/tests/test_array_object.py

Lines changed: 1 addition & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
from .. import ones, arange, reshape, asarray, result_type, all, equal, stack
1212
from .._array_object import Array
13-
from .._devices import CPU_DEVICE, Device, DLDeviceType
13+
from .._devices import CPU_DEVICE, Device
1414
from .._dtypes import (
1515
_all_dtypes,
1616
_boolean_dtypes,
@@ -758,39 +758,6 @@ def test_dlpack_2023_12(api_version):
758758
a.__dlpack__(copy=True)
759759
a.__dlpack__(copy=None)
760760

761-
@pytest.mark.parametrize(
762-
("device", "expected"),
763-
[
764-
(CPU_DEVICE, (DLDeviceType.kDLCPU, 0)),
765-
(Device("device1"), (DLDeviceType.kDLCUDA, 0)),
766-
(Device("device2"), (DLDeviceType.kDLCUDA, 1)),
767-
],
768-
)
769-
def test_dlpack_device_numbers(device, expected):
770-
a = asarray([1, 2, 3], device=device)
771-
assert a.__dlpack_device__() == expected
772-
773-
774-
@pytest.mark.parametrize("device", [CPU_DEVICE, Device("device1"), Device("device2")])
775-
def test_dlpack_export_with_matching_dl_device(device):
776-
if np.lib.NumpyVersion(np.__version__) < "2.1.0":
777-
pytest.skip("dl_device argument requires NumPy >= 2.1.0")
778-
set_array_api_strict_flags(api_version="2023.12")
779-
780-
a = asarray([1, 2, 3], device=device)
781-
a.__dlpack__(dl_device=a.__dlpack_device__())
782-
783-
784-
@pytest.mark.parametrize("device", [Device("device1"), Device("device2")])
785-
def test_dlpack_cross_device_export_buffer_error(device):
786-
if np.lib.NumpyVersion(np.__version__) < "2.1.0":
787-
pytest.skip("dl_device argument requires NumPy >= 2.1.0")
788-
set_array_api_strict_flags(api_version="2023.12")
789-
790-
a = asarray([1, 2, 3], device=device)
791-
with pytest.raises(BufferError):
792-
a.__dlpack__(dl_device=(DLDeviceType.kDLCPU, 0), copy=False)
793-
794761

795762
def test_pickle():
796763
"""Check that arrays are pickleable (despite raising on `__new__`)"""

0 commit comments

Comments
 (0)