Skip to content

Commit 4c360d3

Browse files
committed
Add _is_kdlcpu_device utility to _dlpack.pyx
Simplifies check for device_type = kDLCPU and device_id = 0
1 parent a30b573 commit 4c360d3

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

dpctl/tensor/_dlpack.pyx

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -718,6 +718,11 @@ class _numpy_array_interface_wrapper:
718718
self._memory_owner = memory_owner
719719

720720

721+
cdef bint _is_kdlcpu_device(DLDevice *dev):
722+
"Check if DLTensor.DLDevice denotes (kDLCPU, 0)"
723+
return (dev.device_type == kDLCPU) and (dev.device_id == 0)
724+
725+
721726
cpdef object from_dlpack_capsule(object py_caps):
722727
"""
723728
from_dlpack_capsule(py_caps)
@@ -915,7 +920,7 @@ cpdef object from_dlpack_capsule(object py_caps):
915920
if readonly:
916921
res_ary.flags_ = (res_ary.flags_ & ~USM_ARRAY_WRITABLE)
917922
return res_ary
918-
elif dl_tensor.device.device_type == kDLCPU:
923+
elif _is_kdlcpu_device(&dl_tensor.device):
919924
ary_iface = _numpy_array_interface_from_dl_tensor(dl_tensor, readonly)
920925
if not versioned:
921926
dlm_holder = _DLManagedTensorOwner._create(dlm_tensor)

0 commit comments

Comments
 (0)