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

@@ -157,6 +157,22 @@ unsafe fn blas_1d_params<A>(ptr: *const A, len: usize, stride: isize) -> (*const
 ///
 /// For two-dimensional arrays, the dot method computes the matrix
 /// multiplication.
+///
+/// For higher-dimensional arrays (3-D through 6-D, and dynamic-dimensional),
+/// `Dot<Ix2>` contracts the last axis of the left-hand side with the first
+/// axis of the right-hand side, following NumPy semantics.  For example, if
+/// `self` has shape *I* × *J* × *K* and `rhs` has shape *K* × *N*, the
+/// result has shape *I* × *J* × *N*.
+///
+/// ```
+/// use ndarray::{Array3, Array2};
+/// use ndarray::linalg::Dot;
+///
+/// let a = Array3::<f64>::zeros((3, 4, 5));
+/// let b = Array2::<f64>::zeros((5, 6));
+/// let c = a.dot(&b);
+/// assert_eq!(c.shape(), &[3, 4, 6]);
+/// ```
 pub trait Dot<Rhs: ?Sized>
 {
     /// The result of the operation.
@@ -222,6 +238,112 @@ impl_dots!(Ix1, Ix2);
 impl_dots!(Ix2, Ix1);
 impl_dots!(Ix2, Ix2);
 
+/// Compute the dot product of an N-dimensional LHS array with a 2-D RHS matrix.
+///
+/// Contracts the last axis of `lhs` with the first axis of `rhs`.
+/// If `lhs` has shape *d₀ × d₁ × … × dₙ₋₁ × K* and `rhs` has shape
+/// *K × N*, the result has shape *d₀ × d₁ × … × dₙ₋₁ × N*.
+///
+/// Two paths are used internally:
+/// - **C-contiguous LHS**: the leading axes are flattened into a single
+///   2-D view (zero-copy) and delegated to the optimised 2-D `dot`.
+/// - **Non-contiguous LHS**: a result array is pre-allocated and filled
+///   in-place via recursive `general_mat_mul` calls (no intermediate copies).
+#[track_caller]
+fn nd_dot<A: LinalgScalar>(lhs: &ArrayRef<A, IxDyn>, rhs: &ArrayRef<A, Ix2>) -> Array<A, IxDyn>
+{
+    let ndim = lhs.ndim();
+    let k = lhs.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)",
+            lhs.shape(),
+            rhs.shape()
+        );
+    }
+
+    let mut out_shape = lhs.shape().to_vec();
+    *out_shape.last_mut().unwrap() = n;
+
+    if lhs.is_standard_layout() {
+        // C-contiguous: to_shape returns a *view* (no copy of LHS data).
+        let rows = lhs.len() / k;
+        // unwrap: rows * k == lhs.len() by construction, so reshape always succeeds.
+        let lhs_2d = lhs.to_shape((rows, k)).unwrap();
+        let result_2d = lhs_2d.dot(rhs);
+        // 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: pre-allocate and fill in-place to avoid whole-array copying.
+        let mut out = Array::zeros(IxDyn(&out_shape));
+        nd_dot_non_contiguous(&lhs.view(), &rhs.view(), &mut out.view_mut());
+        out
+    }
+}
+
+/// Recursive helper: writes the N-D × 2-D product directly into `out`
+/// using `general_mat_mul` at the 2-D base case.
+fn nd_dot_non_contiguous<A>(lhs: &ArrayRef<A, IxDyn>, rhs: &ArrayRef<A, Ix2>, out: &mut ArrayRef<A, IxDyn>)
+where A: LinalgScalar
+{
+    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.view().axis_iter(Axis(0)))
+            .and(out.view_mut().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 result_dyn = nd_dot(&self.view().into_dyn(), rhs);
+                // unwrap: nd_dot preserves the number of dimensions.
+                result_dyn.into_dimensionality::<$dim>().unwrap()
+            }
+        }
+
+        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>
+    {
+        nd_dot(&self.view().into_dyn(), rhs)
+    }
+}
+
+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);
+}