Merge pull request #21148 from charris/backport-21132

BUG,ENH: np._from_dlpack: export arrays with any strided size-1 dimensions
numpy · Mar 3, 2022 · 254322c · 254322c
2 parents b49feda + bda7054
commit 254322c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c
@@ -138,7 +138,7 @@ array_dlpack(PyArrayObject *self,
 
     if (!PyArray_IS_C_CONTIGUOUS(self) && PyArray_SIZE(self) != 1) {
         for (int i = 0; i < ndim; ++i) {
-            if (strides[i] % itemsize != 0) {
+            if (shape[i] != 1 && strides[i] % itemsize != 0) {
                 PyErr_SetString(PyExc_RuntimeError,
                         "DLPack only supports strides which are a multiple of "
                         "itemsize.");

diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py
@@ -100,7 +100,7 @@ def dlpack_deleter_exception(self):
         x = np.arange(5)
         _ = x.__dlpack__()
         raise RuntimeError
-    
+
     def test_dlpack_destructor_exception(self):
         with pytest.raises(RuntimeError):
             self.dlpack_deleter_exception()
@@ -110,3 +110,14 @@ def test_readonly(self):
         x.flags.writeable = False
         with pytest.raises(TypeError):
             x.__dlpack__()
+
+    def test_ndim0(self):
+        x = np.array(1.0)
+        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)
+        assert_array_equal(x, y)