Skip to content

Commit fe71e99

Browse files
committed
Prevent losing reference of DLPack holder when returning NumPy array
Passing the used Python capsule to `_numpy_array_interface_wrapper` would cause premature deallocation of the newly created DLPack owner object, causing the dl_tensor to be prematurely deleted Now pass the new owner object instead
1 parent 70c6772 commit fe71e99

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

dpctl/tensor/_dlpack.pyx

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -713,9 +713,9 @@ class _numpy_array_interface_wrapper:
713713
converted to numpy.
714714
"""
715715

716-
def __init__(self, array_interface, pycapsule) -> None:
716+
def __init__(self, array_interface, memory_owner) -> None:
717717
self.__array_interface__ = array_interface
718-
self._pycapsule = pycapsule
718+
self._memory_owner = memory_owner
719719

720720

721721
cpdef object from_dlpack_capsule(object py_caps):
@@ -920,10 +920,11 @@ cpdef object from_dlpack_capsule(object py_caps):
920920
if not versioned:
921921
dlm_holder = _DLManagedTensorOwner._create(dlm_tensor)
922922
cpython.PyCapsule_SetName(py_caps, 'used_dltensor')
923+
return np.ctypeslib.as_array(_numpy_array_interface_wrapper(ary_iface, dlm_holder))
923924
else:
924925
dlmv_holder = _DLManagedTensorVersionedOwner._create(dlmv_tensor)
925926
cpython.PyCapsule_SetName(py_caps, 'used_dltensor_versioned')
926-
return np.ctypeslib.as_array(_numpy_array_interface_wrapper(ary_iface, py_caps))
927+
return np.ctypeslib.as_array(_numpy_array_interface_wrapper(ary_iface, dlmv_holder))
927928
else:
928929
raise BufferError(
929930
"The DLPack tensor resides on unsupported device."

0 commit comments

Comments
 (0)