implement dot for higher-dimensional arrays with Ix2

src/indexes.rs

@@ -57,10 +57,7 @@ where D: Dimension
     #[inline]
     fn next(&mut self) -> Option<Self::Item>
     {
-        let index = match self.index {
-            None => return None,
-            Some(ref ix) => ix.clone(),
-        };
+        let index = self.index.as_ref()?.clone();
         self.index = self.dim.next_for(index.clone());
         Some(index.into_pattern())
     }

src/iterators/mod.rs

@@ -512,10 +512,7 @@ impl<'a, A, D: Dimension> Iterator for IndexedIter<'a, A, D>
     #[inline]
     fn next(&mut self) -> Option<Self::Item>
     {
-        let index = match self.0.inner.index {
-            None => return None,
-            Some(ref ix) => ix.clone(),
-        };
+        let index = self.0.inner.index.as_ref()?.clone();
         match self.0.next() {
             None => None,
             Some(elem) => Some((index.into_pattern(), elem)),
@@ -695,10 +692,7 @@ impl<'a, A, D: Dimension> Iterator for IndexedIterMut<'a, A, D>
     #[inline]
     fn next(&mut self) -> Option<Self::Item>
     {
-        let index = match self.0.inner.index {
-            None => return None,
-            Some(ref ix) => ix.clone(),
-        };
+        let index = self.0.inner.index.as_ref()?.clone();
         match self.0.next() {
             None => None,
             Some(elem) => Some((index.into_pattern(), elem)),

src/linalg/impl_linalg.rs

@@ -222,6 +222,85 @@ impl_dots!(Ix1, Ix2);
 impl_dots!(Ix2, Ix1);
 impl_dots!(Ix2, Ix2);
 
+macro_rules! impl_dot_nd_ix2 {
+    ($dim:ty) => {
+        impl<A> Dot<ArrayRef<A, Ix2>> for ArrayRef<A, $dim>
+        where A: LinalgScalar
+        {
+            type Output = Array<A, $dim>;
+
+            #[track_caller]
+            fn dot(&self, rhs: &ArrayRef<A, Ix2>) -> Array<A, $dim>
+            {
+                let ndim = self.ndim();
+                let k = self.shape()[ndim - 1];
+                let k2 = rhs.shape()[0];
+                let n = rhs.shape()[1];
+                if k != k2 {
+                    dot_shape_error(self.len() / k, k, k2, n);
+                }
+                let rows = self.len() / k;
+                let lhs_2d = self
+                    .to_shape((rows, k))
+                    .expect("ndarray: to_shape failed in nd dot");
+                let result_2d = lhs_2d.dot(rhs);
+
+                let mut out_dim = <$dim>::zeros(ndim);
+                for i in 0..ndim - 1 {
+                    out_dim[i] = self.shape()[i];
+                }
+                out_dim[ndim - 1] = n;
+
+                result_2d
+                    .to_shape(out_dim)
+                    .expect("ndarray: to_shape failed reshaping nd dot result")
+                    .into_owned()
+            }
+        }
+
+        impl_dots!($dim, Ix2);
+    };
+}
+
+impl_dot_nd_ix2!(Ix3);
+impl_dot_nd_ix2!(Ix4);
+impl_dot_nd_ix2!(Ix5);
+impl_dot_nd_ix2!(Ix6);
+
+impl<A> Dot<ArrayRef<A, Ix2>> for ArrayRef<A, IxDyn>
+where A: LinalgScalar
+{
+    type Output = Array<A, IxDyn>;
+
+    #[track_caller]
+    fn dot(&self, rhs: &ArrayRef<A, Ix2>) -> Array<A, IxDyn>
+    {
+        let ndim = self.ndim();
+        let k = self.shape()[ndim - 1];
+        let k2 = rhs.shape()[0];
+        let n = rhs.shape()[1];
+        if k != k2 {
+            dot_shape_error(self.len() / k, k, k2, n);
+        }
+        let rows = self.len() / k;
+        let lhs_2d = self
+            .to_shape((rows, k))
+            .expect("ndarray: to_shape failed in nd dot (IxDyn)");
+        let result_2d = lhs_2d.dot(rhs);
+
+        let mut out_shape = self.shape().to_vec();
+        *out_shape.last_mut().unwrap() = n;
+
+        result_2d
+            .to_shape(IxDyn(&out_shape))
+            .expect("ndarray: to_shape failed reshaping nd dot result (IxDyn)")
+            .into_owned()
+    }
+}
+
+impl_dots!(IxDyn, Ix2);
+impl_dots!(IxDyn, IxDyn);
+
 impl<A> Dot<ArrayRef<A, Ix1>> for ArrayRef<A, Ix1>
 where A: LinalgScalar
 {

tests/oper.rs

@@ -1,6 +1,8 @@
 #![allow(clippy::many_single_char_names, clippy::deref_addrof, clippy::unreadable_literal)]
+#![recursion_limit = "256"]
 use ndarray::linalg::general_mat_mul;
 use ndarray::linalg::kron;
+use ndarray::linalg::Dot;
 use ndarray::prelude::*;
 #[cfg(feature = "approx")]
 use ndarray::Order;
@@ -869,3 +871,158 @@ fn kron_i64()
     let r = arr2(&[[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]);
     assert_eq!(kron(&a, &b), r);
 }
+
+// Helper for the higher-dimensional dot tests below: performs the same
+// operation as ndarray's nd dot but using only the 2-D path, so we can
+// cross-check the results independently.
+fn reference_nd_dot<A, D>(lhs: &Array<A, D>, rhs: &Array2<A>) -> Array<A, D>
+where
+    A: LinalgScalar + std::fmt::Debug,
+    D: Dimension,
+{
+    let ndim = lhs.ndim();
+    let k = lhs.shape()[ndim - 1];
+    let rows = lhs.len() / k;
+    let n = rhs.shape()[1];
+
+    let lhs_2d = lhs.to_shape((rows, k)).unwrap().into_owned();
+    let res_2d = reference_mat_mul(&lhs_2d, rhs);
+
+    let mut out_dim = D::zeros(ndim);
+    for i in 0..ndim - 1 {
+        out_dim[i] = lhs.shape()[i];
+    }
+    out_dim[ndim - 1] = n;
+
+    res_2d.to_shape(out_dim).unwrap().into_owned()
+}
+
+#[test]
+fn dot_3d_by_2d()
+{
+    let lhs: Array3<f64> = ArrayBuilder::new((3, 4, 5)).build();
+    let rhs: Array2<f64> = ArrayBuilder::new((5, 6)).build();
+
+    let result = lhs.dot(&rhs);
+    let expected = reference_nd_dot(&lhs, &rhs);
+
+    assert_eq!(result.shape(), &[3, 4, 6]);
+    assert_eq!(result, expected);
+}
+
+#[test]
+fn dot_3d_by_2d_non_contiguous()
+{
+    // Slice with stride 2 to get a non-contiguous layout.
+    let base: Array3<f64> = ArrayBuilder::new((6, 4, 5)).build();
+    let lhs = base.slice(s![..;2, .., ..]).to_owned();
+    let rhs: Array2<f64> = ArrayBuilder::new((5, 7)).build();
+
+    let result = lhs.dot(&rhs);
+    let expected = reference_nd_dot(&lhs, &rhs);
+
+    assert_eq!(result.shape(), &[3, 4, 7]);
+    assert_eq!(result, expected);
+}
+
+#[test]
+fn dot_3d_by_2d_integer()
+{
+    let lhs: Array3<i32> = ArrayBuilder::new((2, 3, 4)).build();
+    let rhs: Array2<i32> = ArrayBuilder::new((4, 5)).build();
+
+    let result = lhs.dot(&rhs);
+    let expected = reference_nd_dot(&lhs, &rhs);
+
+    assert_eq!(result.shape(), &[2, 3, 5]);
+    assert_eq!(result, expected);
+}
+
+#[test]
+#[should_panic(expected = "not compatible for matrix multiplication")]
+fn dot_3d_by_2d_shape_mismatch()
+{
+    let lhs: Array3<f64> = Array3::zeros((3, 4, 5));
+    let rhs: Array2<f64> = Array2::zeros((6, 7));
+    let _ = lhs.dot(&rhs);
+}
+
+#[test]
+fn dot_4d_by_2d()
+{
+    let lhs: Array4<f64> = ArrayBuilder::new((2, 3, 4, 5)).build();
+    let rhs: Array2<f64> = ArrayBuilder::new((5, 6)).build();
+
+    let result = lhs.dot(&rhs);
+    let expected = reference_nd_dot(&lhs, &rhs);
+
+    assert_eq!(result.shape(), &[2, 3, 4, 6]);
+    assert_eq!(result, expected);
+}
+
+// The shapes here match the NumPy example from issue #1587.
+#[test]
+fn dot_5d_by_2d()
+{
+    let lhs: Array5<f64> =
+        Array5::from_shape_vec((3, 2, 5, 9, 12), (0..3 * 2 * 5 * 9 * 12).map(|x| x as f64).collect()).unwrap();
+    let rhs: Array2<f64> = Array2::from_shape_vec((12, 13), (0..12 * 13).map(|x| x as f64).collect()).unwrap();
+
+    let result = lhs.dot(&rhs);
+    let expected = reference_nd_dot(&lhs, &rhs);
+
+    assert_eq!(result.shape(), &[3, 2, 5, 9, 13]);
+    assert_eq!(result, expected);
+}
+
+#[test]
+fn dot_6d_by_2d()
+{
+    let lhs: Array6<f64> =
+        Array6::from_shape_vec((2, 2, 2, 2, 2, 3), (0..2usize.pow(5) * 3).map(|x| x as f64).collect()).unwrap();
+    let rhs: Array2<f64> = Array2::from_shape_vec((3, 4), (0..12).map(|x| x as f64).collect()).unwrap();
+
+    let result = lhs.dot(&rhs);
+    let expected = reference_nd_dot(&lhs, &rhs);
+
+    assert_eq!(result.shape(), &[2, 2, 2, 2, 2, 4]);
+    assert_eq!(result, expected);
+}
+
+#[test]
+fn dot_dyn_3d_by_2d()
+{
+    let lhs: ArrayD<f64> = ArrayD::from_shape_vec(IxDyn(&[3, 4, 5]), (0..60).map(|x| x as f64).collect()).unwrap();
+    let rhs: Array2<f64> = ArrayBuilder::new((5, 6)).build();
+
+    let result = lhs.dot(&rhs);
+    assert_eq!(result.shape(), &[3, 4, 6]);
+
+    let lhs_fixed: Array3<f64> = lhs.into_dimensionality::<Ix3>().unwrap();
+    let expected = lhs_fixed.dot(&rhs);
+    assert_eq!(result.into_dimensionality::<Ix3>().unwrap(), expected);
+}
+
+#[test]
+fn dot_dyn_5d_by_2d()
+{
+    let lhs: ArrayD<f64> =
+        ArrayD::from_shape_vec(IxDyn(&[3, 2, 5, 9, 12]), (0..3 * 2 * 5 * 9 * 12).map(|x| x as f64).collect()).unwrap();
+    let rhs: Array2<f64> = Array2::from_shape_vec((12, 13), (0..12 * 13).map(|x| x as f64).collect()).unwrap();
+
+    let result = lhs.dot(&rhs);
+    assert_eq!(result.shape(), &[3, 2, 5, 9, 13]);
+
+    let lhs_fixed: Array5<f64> = lhs.into_dimensionality::<Ix5>().unwrap();
+    let expected: Array5<f64> = lhs_fixed.dot(&rhs);
+    assert_eq!(result.into_dimensionality::<Ix5>().unwrap(), expected);
+}
+
+#[test]
+#[should_panic(expected = "not compatible for matrix multiplication")]
+fn dot_dyn_shape_mismatch()
+{
+    let lhs: ArrayD<f64> = ArrayD::zeros(IxDyn(&[3, 4, 5]));
+    let rhs: Array2<f64> = Array2::zeros((6, 7));
+    let _ = lhs.dot(&rhs);
+}