diff --git a/doc/neps/nep-0047-array-api-standard.rst b/doc/neps/nep-0047-array-api-standard.rst index 53b8e35b001f..a94b3b42345b 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 000000000000..75fa9e1817d6 --- /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 5b7c2a5e683f..4301fe553a9f 100644 --- a/doc/source/conf.py +++ b/doc/source/conf.py @@ -307,6 +307,7 @@ def setup(app): '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 25cfa3de1136..33dd16c4ee9e 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 ` instead. + code. Use the :py:doc:`buffer protocol ` 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 30780c286c41..9d2954f2c35b 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 adad4dab92dd..853f324bad8d 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 ` 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 ` 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 `__ 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 `. + +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 "", line 1, in + 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 "", line 1, in + 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 "", line 1, in + 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 297c482e54bf..59a0c6673317 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 741498ff610f..3b014d37b2d6 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 6ac9951fb51c..80409669d174 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -1573,17 +1573,38 @@ 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 - `_ + >>> 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 f88d75978697..1a37ed3e143e 100644 --- a/numpy/core/multiarray.py +++ b/numpy/core/multiarray.py @@ -14,7 +14,7 @@ # 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 @@ '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 @@ 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 2c5265709e8d..fa5e9c67cfcf 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -13,7 +13,7 @@ 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 @@ '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 14ca352c01a7..cb926a26271d 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 e4886cf4bfae..d5b1af101bab 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 12923a6c6658..f206ce2c612f 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 43a663f66206..717210b5423a 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -27,12 +27,12 @@ def test_strides_not_multiple_of_itemsize(self): 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 @@ def test_from_dlpack_refcount(self): ]) 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 @@ def test_invalid_dtype(self): 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 @@ def test_readonly(self): 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 619d6c644d73..eb9a27ab5d44 100644 --- a/tools/refguide_check.py +++ b/tools/refguide_check.py @@ -104,6 +104,8 @@ '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,