Properly handle scalar tensors

Author: Luthaf

src/data_types.rs

@@ -174,6 +174,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_dlpack_pointer_cast() {
         let value: u32 = 42;
         let mut mock_data = value;

src/lib.rs

@@ -475,9 +475,13 @@ impl<'a> DLPackTensorRef<'a> {
 
     /// Get the shape of this tensor
     pub fn shape(&self) -> &[i64] {
-        assert!(!self.raw.shape.is_null());
-        unsafe {
-            return std::slice::from_raw_parts(self.raw.shape, self.n_dims());
+        if self.raw.shape.is_null() {
+            assert!(self.raw.ndim == 0, "Shape pointer is null but ndim is not 0");
+            return &[];
+        } else {
+            unsafe {
+                return std::slice::from_raw_parts(self.raw.shape, self.n_dims());
+            }
         }
     }
 

src/ndarray/mod.rs

@@ -423,6 +423,7 @@ where
             ctx.array.as_ptr()
         };
 
+        let ndim = ctx.shape.len() as i32;
         let dl_tensor = sys::DLTensor {
             // Casting to a mut pointer is not necessarily safe, but is required
             // by DLPack. The data can be mutated through this pointer, we
@@ -433,10 +434,10 @@ where
                 device_type: sys::DLDeviceType::kDLCPU,
                 device_id: 0,
             },
-            ndim: ctx.shape.len() as i32,
+            ndim: ndim,
             dtype: T::get_dlpack_data_type(),
-            shape: ctx.shape.as_mut_ptr(),
-            strides: ctx.strides.as_mut_ptr(),
+            shape: if ndim == 0 { std::ptr::null_mut() } else { ctx.shape.as_mut_ptr() },
+            strides: if ndim == 0 { std::ptr::null_mut() } else { ctx.strides.as_mut_ptr() },
             byte_offset: 0,
         };
 
@@ -489,8 +490,8 @@ where
             },
             ndim,
             dtype: T::get_dlpack_data_type(),
-            shape: ctx.shape.as_mut_ptr(),
-            strides: ctx.strides.as_mut_ptr(),
+            shape: if ndim == 0 { std::ptr::null_mut() } else { ctx.shape.as_mut_ptr() },
+            strides: if ndim == 0 { std::ptr::null_mut() } else { ctx.strides.as_mut_ptr() },
             byte_offset: 0,
         };
 
@@ -786,7 +787,7 @@ mod tests {
         assert_eq!(array_view.shape(), &[0, 0, 0]);
     }
 
-    unsafe extern "C" fn empty_deleter(tensor: *mut sys::DLManagedTensorVersioned) {
+    unsafe extern "C" fn box_deleter(tensor: *mut sys::DLManagedTensorVersioned) {
         let _ = Box::from_raw(tensor);
     }
 
@@ -811,7 +812,7 @@ mod tests {
         let managed = Box::new(crate::sys::DLManagedTensorVersioned {
             version: crate::sys::DLPackVersion::current(),
             manager_ctx: std::ptr::null_mut(),
-            deleter: Some(empty_deleter),
+            deleter: Some(box_deleter),
             flags: 0,
             dl_tensor,
         });
@@ -820,4 +821,79 @@ mod tests {
         let array: Array3<f32> = tensor.try_into().unwrap();
         assert_eq!(array.shape(), &[0, 0, 0]);
     }
+
+    #[test]
+    fn scalar_ndarray_to_dlpack() {
+        let array = arr0(42.0f64);
+        let tensor: DLPackTensor = array.try_into().unwrap();
+        assert_eq!(tensor.n_dims(), 0);
+        assert!(tensor.as_dltensor().shape.is_null());
+        assert!(tensor.shape().is_empty());
+        assert!(tensor.as_dltensor().strides.is_null());
+        assert!(tensor.strides().is_none());
+    }
+
+    #[test]
+    fn scalar_arc_array_to_dlpack() {
+        let array = ndarray::ArcArray::<f64, ndarray::Ix0>::from_elem((), 42.0f64);
+        let tensor: DLPackTensor = array.try_into().unwrap();
+        assert_eq!(tensor.n_dims(), 0);
+        assert!(tensor.as_dltensor().shape.is_null());
+        assert!(tensor.shape().is_empty());
+        assert!(tensor.as_dltensor().strides.is_null());
+        assert!(tensor.strides().is_none());
+    }
+
+    #[test]
+    fn scalar_dlpack_to_ndarray_view() {
+        let mut value = 3.41f32;
+
+        let dl_tensor = DLTensor {
+            data: (&mut value as *mut f32).cast(),
+            device: DLDevice {
+                device_type: DLDeviceType::kDLCPU,
+                device_id: 0,
+            },
+            ndim: 0,
+            dtype: f32::get_dlpack_data_type(),
+            shape: std::ptr::null_mut(),
+            strides: std::ptr::null_mut(),
+            byte_offset: 0,
+        };
+
+        let dlpack_ref = unsafe { DLPackTensorRef::from_raw(dl_tensor) };
+        let array_view = ArrayView0::<f32>::try_from(dlpack_ref).unwrap();
+        assert!(array_view.shape().is_empty());
+        assert_eq!(array_view[()], 3.41);
+    }
+
+    #[test]
+    fn scalar_dlpack_to_ndarray_owned() {
+        let mut value = 2.72f64;
+
+        let dl_tensor = DLTensor {
+            data: (&mut value as *mut f64).cast(),
+            device: DLDevice {
+                device_type: DLDeviceType::kDLCPU,
+                device_id: 0,
+            },
+            ndim: 0,
+            dtype: f64::get_dlpack_data_type(),
+            shape: std::ptr::null_mut(),
+            strides: std::ptr::null_mut(),
+            byte_offset: 0,
+        };
+
+        let managed = Box::new(crate::sys::DLManagedTensorVersioned {
+            version: crate::sys::DLPackVersion::current(),
+            manager_ctx: std::ptr::null_mut(),
+            deleter: Some(box_deleter),
+            flags: 0,
+            dl_tensor,
+        });
+
+        let tensor = unsafe { DLPackTensor::from_ptr(Box::into_raw(managed)) };
+        let array: Array0<f64> = tensor.try_into().unwrap();
+        assert_eq!(array[()], 2.72);
+    }
 }

src/ndarray/sync.rs

@@ -412,4 +412,37 @@ mod tests {
         assert_eq!(tensor.shape(), &[0, 0, 0]);
         assert!(tensor.as_dltensor().data.is_null());
     }
+
+    #[test]
+    fn scalar_mutex_to_dlpack() {
+        let array = Arc::new(Mutex::new(ndarray::arr0(42.0f64)));
+        let tensor: DLPackTensor = Arc::clone(&array).try_into().unwrap();
+        assert_eq!(tensor.n_dims(), 0);
+        assert!(tensor.shape().is_empty());
+        unsafe {
+            assert_eq!(std::ptr::read(tensor.data_ptr::<f64>().unwrap()), 42.0);
+        }
+    }
+
+    #[test]
+    fn scalar_rwlock_write_to_dlpack() {
+        let array = Arc::new(RwLock::new(ndarray::arr0(3.41f64)));
+        let tensor: DLPackTensor = ReadWrite(Arc::clone(&array)).try_into().unwrap();
+        assert_eq!(tensor.n_dims(), 0);
+        assert!(tensor.shape().is_empty());
+        unsafe {
+            assert_eq!(std::ptr::read(tensor.data_ptr::<f64>().unwrap()), 3.41);
+        }
+    }
+
+    #[test]
+    fn scalar_rwlock_read_to_dlpack() {
+        let array = Arc::new(RwLock::new(ndarray::arr0(2.72f64)));
+        let tensor: DLPackTensor = ReadOnly(Arc::clone(&array)).try_into().unwrap();
+        assert_eq!(tensor.n_dims(), 0);
+        assert!(tensor.shape().is_empty());
+        unsafe {
+            assert_eq!(std::ptr::read(tensor.data_ptr::<f64>().unwrap()), 2.72);
+        }
+    }
 }

src/pyo3.rs

@@ -522,6 +522,58 @@ result_capsule = array.__dlpack__()
         })
     }
 
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_scalar_ndarray_to_numpy() -> PyResult<()> {
+        Python::initialize();
+        Python::attach(|py| {
+            let rust_array = ndarray::arr0(3.41f64);
+            let dl_tensor = DLPackTensor::try_from(rust_array).unwrap();
+            let tensor = PyDLPack::try_from(dl_tensor).unwrap();
+
+            let locals = PyDict::new(py);
+            locals.set_item("np", py.import("numpy")?)?;
+            locals.set_item("tensor", tensor)?;
+
+            let code = c_str!(
+                "
+array = np.from_dlpack(tensor)
+assert array.ndim == 0, f\"got ndim {array.ndim}\"
+assert array.size == 1
+assert np.allclose(array, 3.41)
+"
+            );
+            py.run(code, None, Some(&locals))?;
+            Ok(())
+        })
+    }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_scalar_numpy_to_ndarray() -> PyResult<()> {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            locals.set_item("np", py.import("numpy")?)?;
+
+            let code = c_str!(
+                "
+array = np.array(42.0)
+result_capsule = array.__dlpack__()
+"
+            );
+            py.run(code, None, Some(&locals))?;
+
+            let result = locals.get_item("result_capsule")?.unwrap();
+            let capsule: Bound<PyCapsule> = result.extract()?;
+
+            let dlpack_ref = DLPackTensorRef::try_from(capsule)?;
+            let array = ndarray::ArrayView0::<f64>::try_from(dlpack_ref).unwrap();
+            assert_eq!(array[()], 42.0);
+            Ok(())
+        })
+    }
+
     #[test]
     #[cfg_attr(miri, ignore)]
     fn test_v1_0_null_strides_allowed() -> PyResult<()> {

src/sync.rs

@@ -395,4 +395,35 @@ mod tests {
         assert_eq!(tensor.shape(), &[0]);
         assert!(tensor.as_dltensor().data.is_null());
     }
+
+    #[test]
+    fn scalar_mutex_to_dlpack() {
+        let data = Arc::new(Mutex::new(vec![42i32]));
+        let tensor: DLPackTensor = Arc::clone(&data).try_into().unwrap();
+        assert_eq!(tensor.shape(), &[1]);
+
+        unsafe {
+            assert_eq!(std::ptr::read(tensor.data_ptr::<i32>().unwrap()), 42);
+        }
+    }
+
+    #[test]
+    fn scalar_rwlock_write_to_dlpack() {
+        let data = Arc::new(RwLock::new(vec![42i16]));
+        let tensor: DLPackTensor = ReadWrite(Arc::clone(&data)).try_into().unwrap();
+        assert_eq!(tensor.shape(), &[1]);
+        unsafe {
+            assert_eq!(std::ptr::read(tensor.data_ptr::<i16>().unwrap()), 42);
+        }
+    }
+
+    #[test]
+    fn scalar_rwlock_read_to_dlpack() {
+        let data = Arc::new(RwLock::new(vec![42u64]));
+        let tensor: DLPackTensor = ReadOnly(Arc::clone(&data)).try_into().unwrap();
+        assert_eq!(tensor.shape(), &[1]);
+        unsafe {
+            assert_eq!(std::ptr::read(tensor.data_ptr::<u64>().unwrap()), 42);
+        }
+    }
 }

src/sys.rs

@@ -256,7 +256,7 @@ pub struct DLTensor {
     ///
     /// Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow,
     /// TVM, perhaps others) do not adhere to this 256 byte alignment
-    /// requirement on CPU/CUDA/ROCm, and always use `byte_offset=0`.  This must
+    /// requirement on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must
     /// be fixed (after which this note will be updated); at the moment it is
     /// recommended to not rely on the data pointer being correctly aligned.
     ///
@@ -287,17 +287,15 @@ pub struct DLTensor {
     pub shape: *mut i64,
     /// Strides of the tensor (in number of elements, not bytes).
     ///
-    ///  can not be NULL if ndim != 0, must points to
-    ///  an array of ndim elements that specifies the strides,
-    ///  so consumer can always rely on strides[dim] being valid for 0 <= dim < ndim.
+    /// can not be NULL if ndim != 0, must points to an array of ndim elements
+    /// that specifies the strides, so consumer can always rely on strides[dim]
+    /// being valid for 0 <= dim < ndim.
     ///
-    ///  When ndim == 0, strides can be set to NULL.
+    /// When ndim == 0, strides can be set to NULL.
     ///
-    /// NOTE: Before DLPack v1.2, strides can be NULL to indicate contiguous data.
-    ///       This is not allowed in DLPack v1.2 and later. The rationale
+    /// NOTE: Before DLPack v1.2, strides can be NULL to indicate contiguous
+    ///       data. This is not allowed in DLPack v1.2 and later. The rationale
     ///       is to simplify the consumer handling.
-    ///
-    /// When ndim == 0, strides may represent NULL.
     pub strides: *mut i64,
     /// The offset in bytes to the beginning pointer to data
     pub byte_offset: u64,

src/vec.rs

@@ -304,7 +304,7 @@ mod tests {
         }
     }
 
-    unsafe extern "C" fn empty_deleter(tensor: *mut sys::DLManagedTensorVersioned) {
+    unsafe extern "C" fn box_deleter(tensor: *mut sys::DLManagedTensorVersioned) {
         let _ = Box::from_raw(tensor);
     }
 
@@ -329,7 +329,7 @@ mod tests {
         let managed = Box::new(crate::sys::DLManagedTensorVersioned {
             version: crate::sys::DLPackVersion::current(),
             manager_ctx: std::ptr::null_mut(),
-            deleter: Some(empty_deleter),
+            deleter: Some(box_deleter),
             flags: 0,
             dl_tensor,
         });
@@ -338,4 +338,44 @@ mod tests {
         let vec: Vec<i32> = tensor.try_into().unwrap();
         assert!(vec.is_empty());
     }
+
+    #[test]
+    fn scalar_vec_to_dlpack() {
+        let data = vec![42i32];
+        let tensor: DLPackTensor = data.try_into().unwrap();
+        assert_eq!(tensor.shape(), &[1]);
+        assert!(!tensor.as_dltensor().data.is_null());
+    }
+
+    #[test]
+    fn scalar_dlpack_to_vec() {
+        let mut shape = vec![1i64];
+        let mut strides = vec![1i64];
+        let mut value = 42f64;
+
+        let dl_tensor = crate::sys::DLTensor {
+            data: (&mut value as *mut f64).cast(),
+            device: crate::sys::DLDevice {
+                device_type: crate::sys::DLDeviceType::kDLCPU,
+                device_id: 0,
+            },
+            ndim: 1,
+            dtype: f64::get_dlpack_data_type(),
+            shape: shape.as_mut_ptr(),
+            strides: strides.as_mut_ptr(),
+            byte_offset: 0,
+        };
+
+        let managed = Box::new(crate::sys::DLManagedTensorVersioned {
+            version: crate::sys::DLPackVersion::current(),
+            manager_ctx: std::ptr::null_mut(),
+            deleter: Some(box_deleter),
+            flags: 0,
+            dl_tensor,
+        });
+
+        let tensor = unsafe { DLPackTensor::from_ptr(Box::into_raw(managed)) };
+        let vec: Vec<f64> = tensor.try_into().unwrap();
+        assert_eq!(vec, vec![42.0]);
+    }
 }