implement dot for higher-dimensional arrays with Ix2

benches/bench1.rs

@@ -6,6 +6,7 @@ extern crate test;
 
 use std::mem::MaybeUninit;
 
+use ndarray::linalg::Dot;
 use ndarray::{arr0, arr1, arr2, azip, s};
 use ndarray::{Array, Array1, Array2, Axis, Ix, Zip};
 use ndarray::{Array3, Array4, ShapeBuilder};
@@ -908,6 +909,23 @@ fn equality_f32_mixorder(bench: &mut test::Bencher)
     bench.iter(|| a == b);
 }
 
+#[bench]
+fn dot_3d_f64_contiguous(bench: &mut test::Bencher)
+{
+    let a: Array3<f64> = Array::zeros((32, 32, 32));
+    let b: Array2<f64> = Array::zeros((32, 32));
+    bench.iter(|| a.dot(&b));
+}
+
+#[bench]
+fn dot_3d_f64_non_contiguous(bench: &mut test::Bencher)
+{
+    let a_base: Array3<f64> = Array::zeros((64, 32, 32));
+    let a = a_base.slice(s![..;2, .., ..]).to_owned();
+    let b: Array2<f64> = Array::zeros((32, 32));
+    bench.iter(|| a.dot(&b));
+}
+
 #[bench]
 fn dot_f32_16(bench: &mut test::Bencher)
 {

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,139 @@ impl_dots!(Ix1, Ix2);
 impl_dots!(Ix2, Ix1);
 impl_dots!(Ix2, Ix2);
 
+fn nd_dot_non_contiguous<A, S1, S2, S3>(
+    lhs: &ArrayBase<S1, IxDyn>, rhs: &ArrayBase<S2, Ix2>, out: &mut ArrayBase<S3, IxDyn>,
+) where
+    A: LinalgScalar,
+    S1: Data<Elem = A>,
+    S2: Data<Elem = A>,
+    S3: DataMut<Elem = A>,
+{
+    let ndim = lhs.ndim();
+    if ndim == 2 {
+        // unwrap: converting a 2-D ArrayView/ArrayViewMut to Ix2 always succeeds.
+        let lhs_2d = lhs.view().into_dimensionality::<Ix2>().unwrap();
+        let mut out_2d = out.view_mut().into_dimensionality::<Ix2>().unwrap();
+        general_mat_mul(A::one(), &lhs_2d, rhs, A::zero(), &mut out_2d);
+    } else {
+        Zip::from(lhs.axis_iter(Axis(0)))
+            .and(out.axis_iter_mut(Axis(0)))
+            .for_each(|lhs_slice, mut out_slice| {
+                nd_dot_non_contiguous(&lhs_slice, rhs, &mut out_slice);
+            });
+    }
+}
+
+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 {
+                    panic!(
+                        "shapes {:?} and {:?} are not compatible for nd dot \
+                         (last axis of lhs must equal first axis of rhs)",
+                        self.shape(),
+                        rhs.shape()
+                    );
+                }
+
+                if self.is_standard_layout() {
+                    // C-contiguous: to_shape returns a *view* (no copy of LHS data).
+                    let rows = self.len() / k;
+                    // unwrap: rows * k == self.len() by construction, so reshape always succeeds.
+                    let lhs_2d = self.to_shape((rows, k)).unwrap();
+                    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;
+
+                    // unwrap: result_2d is a fresh C-contiguous owned array of the correct size.
+                    result_2d.into_shape_with_order(out_dim).unwrap()
+                } else {
+                    // Non-contiguous: iterate over the first axis and write results directly
+                    // into the pre-allocated output array to avoid whole-array copying.
+                    let mut out_dim = <$dim>::zeros(ndim);
+                    for i in 0..ndim - 1 {
+                        out_dim[i] = self.shape()[i];
+                    }
+                    out_dim[ndim - 1] = n;
+                    let mut out = Array::zeros(out_dim);
+                    nd_dot_non_contiguous(&self.view().into_dyn(), &rhs.view(), &mut out.view_mut().into_dyn());
+                    out
+                }
+            }
+        }
+
+        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 {
+            panic!(
+                "shapes {:?} and {:?} are not compatible for nd dot \
+                 (last axis of lhs must equal first axis of rhs)",
+                self.shape(),
+                rhs.shape()
+            );
+        }
+
+        if self.is_standard_layout() {
+            // C-contiguous: to_shape returns a *view* (no copy of LHS data).
+            let rows = self.len() / k;
+            // unwrap: rows * k == self.len() by construction, so reshape always succeeds.
+            let lhs_2d = self.to_shape((rows, k)).unwrap();
+            let result_2d = lhs_2d.dot(rhs);
+
+            let mut out_shape = self.shape().to_vec();
+            *out_shape.last_mut().unwrap() = n;
+
+            // unwrap: result_2d is a fresh C-contiguous owned array of the correct size.
+            result_2d.into_shape_with_order(IxDyn(&out_shape)).unwrap()
+        } else {
+            // Non-contiguous: iterate recursively over the first axis and write results
+            // directly into the pre-allocated output array to avoid whole-array copying.
+            let mut out_shape = self.shape().to_vec();
+            *out_shape.last_mut().unwrap() = n;
+            let mut out = Array::zeros(IxDyn(&out_shape));
+            nd_dot_non_contiguous(&self.view(), &rhs.view(), &mut out.view_mut());
+            out
+        }
+    }
+}
+
+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,174 @@ 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 = "are not compatible for nd dot")]
+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]
+fn dot_dyn_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, .., ..]).into_dyn();
+    let rhs: Array2<f64> = ArrayBuilder::new((5, 7)).build();
+
+    let result = lhs.dot(&rhs);
+    assert_eq!(result.shape(), &[3, 4, 7]);
+
+    let lhs_fixed: ArrayView3<f64> = lhs.into_dimensionality::<Ix3>().unwrap();
+    let expected = lhs_fixed.dot(&rhs);
+    assert_eq!(result.into_dimensionality::<Ix3>().unwrap(), expected);
+}
+
+#[test]
+#[should_panic(expected = "are not compatible for nd dot")]
+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);
+}