BUG,ENH: np._from_dlpack: export arrays with any-strided size-1 dimen…

…sions
numpy · Mar 1, 2022 · 7532ab6 · 7532ab6
1 parent 13e8ace
commit 7532ab6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c
@@ -137,7 +137,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
@@ -115,3 +115,9 @@ 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)