wkcn commented on a change in pull request #18342:
URL: https://github.com/apache/incubator-mxnet/pull/18342#discussion_r426277514
##########
File path: python/mxnet/numpy_extension/utils.py
##########
@@ -119,3 +131,112 @@ def load(file):
return dict(
(py_str(names[i]), ndarray(NDArrayHandle(handles[i])))
for i in range(out_size.value))
+
+
+def from_dlpack(dlpack):
+ """Returns a np.ndarray backed by a dlpack tensor.
+
+ Parameters
+ ----------
+ dlpack: PyCapsule (the pointer of DLManagedTensor)
+ input data
+
+ Returns
+ -------
+ np.ndarray
+ an ndarray backed by a dlpack tensor
+
+ Examples
+ --------
+ >>> x = mx.np.ones((2,3))
+ >>> y = mx.npx.to_dlpack_for_read(x)
+ >>> type(y)
+ <class 'PyCapsule'>
+ >>> z = mx.npx.from_dlpack(y)
+ >>> type(z)
+ <class 'mxnet.numpy.ndarray'>
+ >>> z
+ array([[1., 1., 1.],
+ [1., 1., 1.]])
+
+ >>> w = mx.npx.to_dlpack_for_write(x)
+ >>> type(w)
+ <class 'PyCapsule'>
+ >>> u = mx.npx.from_dlpack(w)
+ >>> u += 1
+ >>> x
+ array([[2., 2., 2.],
+ [2., 2., 2.]])
+ """
+ handle = NDArrayHandle()
+ dlpack = ctypes.py_object(dlpack)
+ assert ctypes.pythonapi.PyCapsule_IsValid(dlpack, _c_str_dltensor),
ValueError(
+ 'Invalid DLPack Tensor. DLTensor capsules can be consumed only once.')
+ dlpack_handle =
ctypes.c_void_p(ctypes.pythonapi.PyCapsule_GetPointer(dlpack, _c_str_dltensor))
+ check_call(_LIB.MXNDArrayFromDLPackEx(dlpack_handle, False,
ctypes.byref(handle)))
+ # Rename PyCapsule (DLPack)
+ ctypes.pythonapi.PyCapsule_SetName(dlpack, _c_str_used_dltensor)
+ # delete the deleter of the old dlpack
+ ctypes.pythonapi.PyCapsule_SetDestructor(dlpack, None)
+ return ndarray(handle=handle)
+
+def to_dlpack_for_read(data):
+ """Returns a reference view of np.ndarray that represents as
DLManagedTensor until
+ all previous write operations on the current array are finished.
+
+ Parameters
+ ----------
+ data: ndarray
Review comment:
data: np.ndarray
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
