Merge pull request #21145 from tirthasheshpatel/fix-gh20743

MAINT, DOC: make np._from_dlpack public

16 files changed, 164 insertions, 32 deletions

diff --git a/doc/neps/nep-0047-array-api-standard.rst b/doc/neps/nep-0047-array-api-standard.rst
index 53b8e35b0..a94b3b423 100644
--- a/doc/neps/nep-0047-array-api-standard.rst
+++ b/doc/neps/nep-0047-array-api-standard.rst
@@ -340,7 +340,7 @@ Adding support for DLPack to NumPy entails:
 
 - Adding a ``ndarray.__dlpack__()`` method which returns a ``dlpack`` C
   structure wrapped in a ``PyCapsule``.
-- Adding a ``np._from_dlpack(obj)`` function, where ``obj`` supports
+- Adding a ``np.from_dlpack(obj)`` function, where ``obj`` supports
   ``__dlpack__()``, and returns an ``ndarray``.
 
 DLPack is currently a ~200 LoC header, and is meant to be included directly, so
diff --git a/doc/release/upcoming_changes/21145.new_function.rst b/doc/release/upcoming_changes/21145.new_function.rst
new file mode 100644
index 000000000..75fa9e181
--- /dev/null
+++ b/doc/release/upcoming_changes/21145.new_function.rst
@@ -0,0 +1,6 @@
+NumPy now supports the DLPack protocol
+--------------------------------------
+`numpy.from_dlpack` has been added to NumPy to exchange data using the DLPack protocol.
+It accepts Python objects that implement the ``__dlpack__`` and ``__dlpack_device__``
+methods and returns a ndarray object which is generally the view of the data of the input
+object.
diff --git a/doc/source/conf.py b/doc/source/conf.py
index 5b7c2a5e6..4301fe553 100644
--- a/doc/source/conf.py
+++ b/doc/source/conf.py
@@ -307,6 +307,7 @@ intersphinx_mapping = {
     'pytest': ('https://docs.pytest.org/en/stable', None),
     'numpy-tutorials': ('https://numpy.org/numpy-tutorials', None),
     'numpydoc': ('https://numpydoc.readthedocs.io/en/latest', None),
+    'dlpack': ('https://dmlc.github.io/dlpack/latest', None)
 }
 
 
diff --git a/doc/source/reference/arrays.interface.rst b/doc/source/reference/arrays.interface.rst
index 25cfa3de1..33dd16c4e 100644
--- a/doc/source/reference/arrays.interface.rst
+++ b/doc/source/reference/arrays.interface.rst
@@ -247,7 +247,8 @@ flag is present.
 .. note::
 
     :obj:`__array_struct__` is considered legacy and should not be used for new
-    code. Use the :py:doc:`buffer protocol <c-api/buffer>` instead.
+    code. Use the :py:doc:`buffer protocol <c-api/buffer>` or the DLPack protocol
+    `numpy.from_dlpack` instead.
 
 
 Type description examples
diff --git a/doc/source/reference/routines.array-creation.rst b/doc/source/reference/routines.array-creation.rst
index 30780c286..9d2954f2c 100644
--- a/doc/source/reference/routines.array-creation.rst
+++ b/doc/source/reference/routines.array-creation.rst
@@ -35,6 +35,7 @@ From existing data
    asmatrix
    copy
    frombuffer
+   from_dlpack
    fromfile
    fromfunction
    fromiter
diff --git a/doc/source/user/basics.interoperability.rst b/doc/source/user/basics.interoperability.rst
index adad4dab9..853f324ba 100644
--- a/doc/source/user/basics.interoperability.rst
+++ b/doc/source/user/basics.interoperability.rst
@@ -55,6 +55,14 @@ describes its memory layout and NumPy does everything else (zero-copy if
 possible). If that's not possible, the object itself is responsible for
 returning a ``ndarray`` from ``__array__()``.
 
+:doc:`DLPack <dlpack:index>` is yet another protocol to convert foriegn objects
+to NumPy arrays in a language and device agnostic manner. NumPy doesn't implicitly
+convert objects to ndarrays using DLPack. It provides the function
+`numpy.from_dlpack` that accepts any object implementing the ``__dlpack__`` method
+and outputs a NumPy ndarray (which is generally a view of the input object's data
+buffer). The :ref:`dlpack:python-spec` page explains the ``__dlpack__`` protocol
+in detail.
+
 The array interface protocol
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -118,6 +126,26 @@ as the original object and any attributes/behavior it may have had, is lost.
 To see an example of a custom array implementation including the use of
 ``__array__()``, see :ref:`basics.dispatch`.
 
+The DLPack Protocol
+~~~~~~~~~~~~~~~~~~~
+
+The :doc:`DLPack <dlpack:index>` protocol defines a memory-layout of
+strided n-dimensional array objects. It offers the following syntax
+for data exchange:
+
+1. A `numpy.from_dlpack` function, which accepts (array) objects with a
+   ``__dlpack__`` method and uses that method to construct a new array
+   containing the data from ``x``.
+2. ``__dlpack__(self, stream=None)`` and ``__dlpack_device__`` methods on the
+   array object, which will be called from within ``from_dlpack``, to query
+   what device the array is on (may be needed to pass in the correct
+   stream, e.g. in the case of multiple GPUs) and to access the data.
+
+Unlike the buffer protocol, DLPack allows exchanging arrays containing data on
+devices other than the CPU (e.g. Vulkan or GPU). Since NumPy only supports CPU,
+it can only convert objects whose data exists on the CPU. But other libraries,
+like PyTorch_ and CuPy_, may exchange data on GPU using this protocol.
+
 
 2. Operating on foreign objects without converting
 --------------------------------------------------
@@ -395,6 +423,78 @@ See `the Dask array documentation
 and the `scope of Dask arrays interoperability with NumPy arrays
 <https://docs.dask.org/en/stable/array.html#scope>`__ for details.
 
+Example: DLPack
+~~~~~~~~~~~~~~~
+
+Several Python data science libraries implement the ``__dlpack__`` protocol.
+Among them are PyTorch_ and CuPy_. A full list of libraries that implement
+this protocol can be found on
+:doc:`this page of DLPack documentation <dlpack:index>`.
+
+Convert a PyTorch CPU tensor to NumPy array:
+
+ >>> import torch
+ >>> x_torch = torch.arange(5)
+ >>> x_torch
+ tensor([0, 1, 2, 3, 4])
+ >>> x_np = np.from_dlpack(x_torch)
+ >>> x_np
+ array([0, 1, 2, 3, 4])
+ >>> # note that x_np is a view of x_torch
+ >>> x_torch[1] = 100
+ >>> x_torch
+ tensor([  0, 100,   2,   3,   4])
+ >>> x_np
+ array([  0, 100,   2,   3,   4])
+
+The imported arrays are read-only so writing or operating in-place will fail:
+
+ >>> x.flags.writeable
+ False
+ >>> x_np[1] = 1
+ Traceback (most recent call last):
+   File "<stdin>", line 1, in <module>
+ ValueError: assignment destination is read-only
+
+A copy must be created in order to operate on the imported arrays in-place, but
+will mean duplicating the memory. Do not do this for very large arrays:
+
+ >>> x_np_copy = x_np.copy()
+ >>> x_np_copy.sort()  # works
+
+.. note::
+
+  Note that GPU tensors can't be converted to NumPy arrays since NumPy doesn't
+  support GPU devices:
+
+   >>> x_torch = torch.arange(5, device='cuda')
+   >>> np.from_dlpack(x_torch)
+   Traceback (most recent call last):
+     File "<stdin>", line 1, in <module>
+   RuntimeError: Unsupported device in DLTensor.
+
+  But, if both libraries support the device the data buffer is on, it is
+  possible to use the ``__dlpack__`` protocol (e.g. PyTorch_ and CuPy_):
+
+   >>> x_torch = torch.arange(5, device='cuda')
+   >>> x_cupy = cupy.from_dlpack(x_torch)
+
+Similarly, a NumPy array can be converted to a PyTorch tensor:
+
+ >>> x_np = np.arange(5)
+ >>> x_torch = torch.from_dlpack(x_np)
+
+Read-only arrays cannot be exported:
+
+ >>> x_np = np.arange(5)
+ >>> x_np.flags.writeable = False
+ >>> torch.from_dlpack(x_np)  # doctest: +ELLIPSIS
+ Traceback (most recent call last):
+   File "<stdin>", line 1, in <module>
+   File ".../site-packages/torch/utils/dlpack.py", line 63, in from_dlpack
+     dlpack = ext_tensor.__dlpack__()
+ TypeError: NumPy currently only supports dlpack for writeable arrays
+
 Further reading
 ---------------
 
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 297c482e5..59a0c6673 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -4372,4 +4372,4 @@ class chararray(ndarray[_ShapeType, _CharDType]):
 class _SupportsDLPack(Protocol[_T_contra]):
     def __dlpack__(self, *, stream: None | _T_contra = ...) -> _PyCapsule: ...
 
-def _from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ...
+def from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ...
diff --git a/numpy/array_api/_creation_functions.py b/numpy/array_api/_creation_functions.py
index 741498ff6..3b014d37b 100644
--- a/numpy/array_api/_creation_functions.py
+++ b/numpy/array_api/_creation_functions.py
@@ -154,7 +154,7 @@ def eye(
 def from_dlpack(x: object, /) -> Array:
     from ._array_object import Array
 
-    return Array._new(np._from_dlpack(x))
+    return Array._new(np.from_dlpack(x))
 
 
 def full(
diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py
index 6ac9951fb..80409669d 100644
--- a/numpy/core/_add_newdocs.py
+++ b/numpy/core/_add_newdocs.py
@@ -1573,17 +1573,38 @@ add_newdoc('numpy.core.multiarray', 'frombuffer',
         array_function_like_doc,
     ))
 
-add_newdoc('numpy.core.multiarray', '_from_dlpack',
+add_newdoc('numpy.core.multiarray', 'from_dlpack',
     """
-    _from_dlpack(x, /)
+    from_dlpack(x, /)
 
     Create a NumPy array from an object implementing the ``__dlpack__``
-    protocol.
+    protocol. Generally, the returned NumPy array is a read-only view
+    of the input object. See [1]_ and [2]_ for more details.
 
-    See Also
+    Parameters
+    ----------
+    x : object
+        A Python object that implements the ``__dlpack__`` and
+        ``__dlpack_device__`` methods.
+
+    Returns
+    -------
+    out : ndarray
+
+    References
+    ----------
+    .. [1] Array API documentation,
+       https://data-apis.org/array-api/latest/design_topics/data_interchange.html#syntax-for-data-interchange-with-dlpack
+
+    .. [2] Python specification for DLPack,
+       https://dmlc.github.io/dlpack/latest/python_spec.html
+
+    Examples
     --------
-    `Array API documentation
-    <https://data-apis.org/array-api/latest/design_topics/data_interchange.html#syntax-for-data-interchange-with-dlpack>`_
+    >>> import torch
+    >>> x = torch.arange(10)
+    >>> # create a view of the torch tensor "x" in NumPy
+    >>> y = np.from_dlpack(x)
     """)
 
 add_newdoc('numpy.core', 'fastCopyAndTranspose',
diff --git a/numpy/core/multiarray.py b/numpy/core/multiarray.py
index f88d75978..1a37ed3e1 100644
--- a/numpy/core/multiarray.py
+++ b/numpy/core/multiarray.py
@@ -14,7 +14,7 @@ from ._multiarray_umath import *  # noqa: F403
 # do not change them. issue gh-15518
 # _get_ndarray_c_version is semi-public, on purpose not added to __all__
 from ._multiarray_umath import (
-    _fastCopyAndTranspose, _flagdict, _from_dlpack, _insert, _reconstruct,
+    _fastCopyAndTranspose, _flagdict, from_dlpack, _insert, _reconstruct,
     _vec_string, _ARRAY_API, _monotonicity, _get_ndarray_c_version,
     _set_madvise_hugepage,
     )
@@ -24,7 +24,7 @@ __all__ = [
     'ITEM_HASOBJECT', 'ITEM_IS_POINTER', 'LIST_PICKLE', 'MAXDIMS',
     'MAY_SHARE_BOUNDS', 'MAY_SHARE_EXACT', 'NEEDS_INIT', 'NEEDS_PYAPI',
     'RAISE', 'USE_GETITEM', 'USE_SETITEM', 'WRAP', '_fastCopyAndTranspose',
-    '_flagdict', '_from_dlpack', '_insert', '_reconstruct', '_vec_string',
+    '_flagdict', 'from_dlpack', '_insert', '_reconstruct', '_vec_string',
     '_monotonicity', 'add_docstring', 'arange', 'array', 'asarray',
     'asanyarray', 'ascontiguousarray', 'asfortranarray', 'bincount',
     'broadcast', 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast',
@@ -47,7 +47,7 @@ _reconstruct.__module__ = 'numpy.core.multiarray'
 scalar.__module__ = 'numpy.core.multiarray'
 
 
-_from_dlpack.__module__ = 'numpy'
+from_dlpack.__module__ = 'numpy'
 arange.__module__ = 'numpy'
 array.__module__ = 'numpy'
 asarray.__module__ = 'numpy'
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 5f6602c71..39bfc2d55 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -13,7 +13,7 @@ from .multiarray import (
     WRAP, arange, array, asarray, asanyarray, ascontiguousarray,
     asfortranarray, broadcast, can_cast, compare_chararrays,
     concatenate, copyto, dot, dtype, empty,
-    empty_like, flatiter, frombuffer, _from_dlpack, fromfile, fromiter,
+    empty_like, flatiter, frombuffer, from_dlpack, fromfile, fromiter,
     fromstring, inner, lexsort, matmul, may_share_memory,
     min_scalar_type, ndarray, nditer, nested_iters, promote_types,
     putmask, result_type, set_numeric_ops, shares_memory, vdot, where,
@@ -41,7 +41,7 @@ __all__ = [
     'newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc',
     'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray',
     'asfortranarray', 'zeros', 'count_nonzero', 'empty', 'broadcast', 'dtype',
-    'fromstring', 'fromfile', 'frombuffer', '_from_dlpack', 'where',
+    'fromstring', 'fromfile', 'frombuffer', 'from_dlpack', 'where',
     'argwhere', 'copyto', 'concatenate', 'fastCopyAndTranspose', 'lexsort',
     'set_numeric_ops', 'can_cast', 'promote_types', 'min_scalar_type',
     'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like',
diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h
index 14ca352c0..cb926a262 100644
--- a/numpy/core/src/common/npy_dlpack.h
+++ b/numpy/core/src/common/npy_dlpack.h
@@ -23,6 +23,6 @@ array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args));
 
 
 NPY_NO_EXPORT PyObject *
-_from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj);
+from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj);
 
 #endif
diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c
index e4886cf4b..d5b1af101 100644
--- a/numpy/core/src/multiarray/dlpack.c
+++ b/numpy/core/src/multiarray/dlpack.c
@@ -269,7 +269,7 @@ array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args))
 }
 
 NPY_NO_EXPORT PyObject *
-_from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) {
+from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) {
     PyObject *capsule = PyObject_CallMethod((PyObject *)obj->ob_type,
             "__dlpack__", "O", obj);
     if (capsule == NULL) {
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 12923a6c6..f206ce2c6 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -4485,7 +4485,7 @@ static struct PyMethodDef array_module_methods[] = {
     {"_reload_guard", (PyCFunction)_reload_guard,
         METH_NOARGS,
         "Give a warning on reload and big warning in sub-interpreters."},
-    {"_from_dlpack", (PyCFunction)_from_dlpack,
+    {"from_dlpack", (PyCFunction)from_dlpack,
         METH_O, NULL},
     {NULL, NULL, 0, NULL}                /* sentinel */
 };
diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py
index 43a663f66..717210b54 100644
--- a/numpy/core/tests/test_dlpack.py
+++ b/numpy/core/tests/test_dlpack.py
@@ -27,12 +27,12 @@ class TestDLPack:
         z = y['int']
 
         with pytest.raises(RuntimeError):
-            np._from_dlpack(z)
+            np.from_dlpack(z)
 
     @pytest.mark.skipif(IS_PYPY, reason="PyPy can't get refcounts.")
     def test_from_dlpack_refcount(self):
         x = np.arange(5)
-        y = np._from_dlpack(x)
+        y = np.from_dlpack(x)
         assert sys.getrefcount(x) == 3
         del y
         assert sys.getrefcount(x) == 2
@@ -45,7 +45,7 @@ class TestDLPack:
     ])
     def test_dtype_passthrough(self, dtype):
         x = np.arange(5, dtype=dtype)
-        y = np._from_dlpack(x)
+        y = np.from_dlpack(x)
 
         assert y.dtype == x.dtype
         assert_array_equal(x, y)
@@ -54,44 +54,44 @@ class TestDLPack:
         x = np.asarray(np.datetime64('2021-05-27'))
 
         with pytest.raises(TypeError):
-            np._from_dlpack(x)
+            np.from_dlpack(x)
 
     def test_invalid_byte_swapping(self):
         dt = np.dtype('=i8').newbyteorder()
         x = np.arange(5, dtype=dt)
 
         with pytest.raises(TypeError):
-            np._from_dlpack(x)
+            np.from_dlpack(x)
 
     def test_non_contiguous(self):
         x = np.arange(25).reshape((5, 5))
 
         y1 = x[0]
-        assert_array_equal(y1, np._from_dlpack(y1))
+        assert_array_equal(y1, np.from_dlpack(y1))
 
         y2 = x[:, 0]
-        assert_array_equal(y2, np._from_dlpack(y2))
+        assert_array_equal(y2, np.from_dlpack(y2))
 
         y3 = x[1, :]
-        assert_array_equal(y3, np._from_dlpack(y3))
+        assert_array_equal(y3, np.from_dlpack(y3))
 
         y4 = x[1]
-        assert_array_equal(y4, np._from_dlpack(y4))
+        assert_array_equal(y4, np.from_dlpack(y4))
 
         y5 = np.diagonal(x).copy()
-        assert_array_equal(y5, np._from_dlpack(y5))
+        assert_array_equal(y5, np.from_dlpack(y5))
 
     @pytest.mark.parametrize("ndim", range(33))
     def test_higher_dims(self, ndim):
         shape = (1,) * ndim
         x = np.zeros(shape, dtype=np.float64)
 
-        assert shape == np._from_dlpack(x).shape
+        assert shape == np.from_dlpack(x).shape
 
     def test_dlpack_device(self):
         x = np.arange(5)
         assert x.__dlpack_device__() == (1, 0)
-        y = np._from_dlpack(x)
+        y = np.from_dlpack(x)
         assert y.__dlpack_device__() == (1, 0)
         z = y[::2]
         assert z.__dlpack_device__() == (1, 0)
@@ -113,11 +113,11 @@ class TestDLPack:
 
     def test_ndim0(self):
         x = np.array(1.0)
-        y = np._from_dlpack(x)
+        y = np.from_dlpack(x)
         assert_array_equal(x, y)
 
     def test_size1dims_arrays(self):
         x = np.ndarray(dtype='f8', shape=(10, 5, 1), strides=(8, 80, 4),
                        buffer=np.ones(1000, dtype=np.uint8), order='F')
-        y = np._from_dlpack(x)
+        y = np.from_dlpack(x)
         assert_array_equal(x, y)
diff --git a/tools/refguide_check.py b/tools/refguide_check.py
index 619d6c644..eb9a27ab5 100644
--- a/tools/refguide_check.py
+++ b/tools/refguide_check.py
@@ -104,6 +104,8 @@ DOCTEST_SKIPDICT = {
     'numpy.random.vonmises': None,
     'numpy.random.power': None,
     'numpy.random.zipf': None,
+    # cases where NumPy docstrings import things from other 3'rd party libs:
+    'numpy.core.from_dlpack': None,
     # remote / local file IO with DataSource is problematic in doctest:
     'numpy.lib.DataSource': None,
     'numpy.lib.Repository': None,