From bda705477e28a1ef8496cbc680650ad5fbed0a97 Mon Sep 17 00:00:00 2001 From: Tirth Patel Date: Tue, 1 Mar 2022 00:40:26 +0530 Subject: [PATCH] BUG,ENH: np._from_dlpack: export arrays with any-strided size-1 dimensions --- numpy/core/src/multiarray/dlpack.c | 2 +- numpy/core/tests/test_dlpack.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index 8491ed5b9bde..37b0b6dcfda8 100644 --- 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 index 2ab55e903861..eb9a1765423c 100644 --- 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)