Extend the usage of the ArrayOrScalar trait to the einsum macro and function.

Author: adamreichold

CHANGELOG.md

@@ -1,7 +1,7 @@
 # Changelog
 
 - Unreleased
-  - The `inner` and `dot` functions can also return a scalar instead of a zero-dimensional array to match NumPy's types ([#285](https://github.com/PyO3/rust-numpy/pull/285))
+  - The `inner`, `dot` and `einsum` functions can also return a scalar instead of a zero-dimensional array to match NumPy's types ([#285](https://github.com/PyO3/rust-numpy/pull/285))
 
 - v0.16.1
   - Fix build when PyO3's `multiple-pymethods` feature is used. ([#288](https://github.com/PyO3/rust-numpy/pull/288))

src/lib.rs

@@ -64,7 +64,7 @@ pub use crate::readonly::{
     PyReadonlyArray, PyReadonlyArray1, PyReadonlyArray2, PyReadonlyArray3, PyReadonlyArray4,
     PyReadonlyArray5, PyReadonlyArray6, PyReadonlyArrayDyn,
 };
-pub use crate::sum_products::{dot, einsum_impl, inner};
+pub use crate::sum_products::{dot, einsum, inner};
 pub use ndarray::{array, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
 
 #[cfg(doctest)]

src/sum_products.rs

@@ -133,13 +133,10 @@ where
 /// Return the Einstein summation convention of given tensors.
 ///
 /// This is usually invoked via the the [`einsum!`] macro.
-pub fn einsum_impl<'py, T, DOUT>(
-    subscripts: &str,
-    arrays: &[&'py PyArray<T, IxDyn>],
-) -> PyResult<&'py PyArray<T, DOUT>>
+pub fn einsum<'py, T, OUT>(subscripts: &str, arrays: &[&'py PyArray<T, IxDyn>]) -> PyResult<OUT>
 where
-    DOUT: Dimension,
     T: Element,
+    OUT: ArrayOrScalar<'py, T>,
 {
     let subscripts = match CStr::from_bytes_with_nul(subscripts.as_bytes()) {
         Ok(subscripts) => Cow::Borrowed(subscripts),
@@ -173,13 +170,13 @@ where
 /// ```
 /// use pyo3::Python;
 /// use ndarray::array;
-/// use numpy::{einsum, pyarray, PyArray};
+/// use numpy::{einsum, pyarray, PyArray, PyArray2};
 ///
 /// Python::with_gil(|py| {
 ///     let tensor = PyArray::arange(py, 0, 2 * 3 * 4, 1).reshape([2, 3, 4]).unwrap();
 ///     let another_tensor = pyarray![py, [20, 30], [40, 50], [60, 70]];
 ///
-///     let result = einsum!("ijk,ji->ik", tensor, another_tensor).unwrap();
+///     let result: &PyArray2<_> = einsum!("ijk,ji->ik", tensor, another_tensor).unwrap();
 ///
 ///     assert_eq!(
 ///         result.readonly().as_array(),
@@ -193,6 +190,6 @@ where
 macro_rules! einsum {
     ($subscripts:literal $(,$array:ident)+ $(,)*) => {{
         let arrays = [$($array.to_dyn(),)+];
-        $crate::einsum_impl(concat!($subscripts, "\0"), &arrays)
+        $crate::einsum(concat!($subscripts, "\0"), &arrays)
     }};
 }

tests/sum_products.rs

@@ -62,22 +62,22 @@ fn test_einsum() {
         let b = pyarray![py, 0, 1, 2, 3, 4];
         let c = pyarray![py, [0, 1, 2], [3, 4, 5]];
 
-        assert_eq!(einsum!("ii", a).unwrap().item(), 60);
-        assert_eq!(
-            einsum!("ii->i", a).unwrap().readonly().as_array(),
-            array![0, 6, 12, 18, 24],
-        );
-        assert_eq!(
-            einsum!("ij->i", a).unwrap().readonly().as_array(),
-            array![10, 35, 60, 85, 110],
-        );
-        assert_eq!(
-            einsum!("ji", c).unwrap().readonly().as_array(),
-            array![[0, 3], [1, 4], [2, 5]],
-        );
-        assert_eq!(
-            einsum!("ij,j", a, b).unwrap().readonly().as_array(),
-            array![30, 80, 130, 180, 230],
-        );
+        let d: &PyArray0<_> = einsum!("ii", a).unwrap();
+        assert_eq!(d.item(), 60);
+
+        let d: i32 = einsum!("ii", a).unwrap();
+        assert_eq!(d, 60);
+
+        let d: &PyArray1<_> = einsum!("ii->i", a).unwrap();
+        assert_eq!(d.readonly().as_array(), array![0, 6, 12, 18, 24]);
+
+        let d: &PyArray1<_> = einsum!("ij->i", a).unwrap();
+        assert_eq!(d.readonly().as_array(), array![10, 35, 60, 85, 110]);
+
+        let d: &PyArray2<_> = einsum!("ji", c).unwrap();
+        assert_eq!(d.readonly().as_array(), array![[0, 3], [1, 4], [2, 5]]);
+
+        let d: &PyArray1<_> = einsum!("ij,j", a, b).unwrap();
+        assert_eq!(d.readonly().as_array(), array![30, 80, 130, 180, 230]);
     });
 }