diff --git a/benches/bench1.rs b/benches/bench1.rs index 3b5405329..6c54a9456 100644 --- a/benches/bench1.rs +++ b/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 = Array::zeros((32, 32, 32)); + let b: Array2 = Array::zeros((32, 32)); + bench.iter(|| a.dot(&b)); +} + +#[bench] +fn dot_3d_f64_non_contiguous(bench: &mut test::Bencher) +{ + let a_base: Array3 = Array::zeros((64, 32, 32)); + let a = a_base.slice(s![..;2, .., ..]).to_owned(); + let b: Array2 = Array::zeros((32, 32)); + bench.iter(|| a.dot(&b)); +} + #[bench] fn dot_f32_16(bench: &mut test::Bencher) { diff --git a/src/indexes.rs b/src/indexes.rs index 0fa2b50fb..b661d3b46 100644 --- a/src/indexes.rs +++ b/src/indexes.rs @@ -57,10 +57,7 @@ where D: Dimension #[inline] fn next(&mut self) -> Option { - 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()) } diff --git a/src/iterators/mod.rs b/src/iterators/mod.rs index abca3579d..fe0a4775b 100644 --- a/src/iterators/mod.rs +++ b/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 { - 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 { - 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)), diff --git a/src/linalg/impl_linalg.rs b/src/linalg/impl_linalg.rs index 81c942bc3..207b9f0eb 100644 --- a/src/linalg/impl_linalg.rs +++ b/src/linalg/impl_linalg.rs @@ -157,6 +157,22 @@ unsafe fn blas_1d_params(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` 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::::zeros((3, 4, 5)); +/// let b = Array2::::zeros((5, 6)); +/// let c = a.dot(&b); +/// assert_eq!(c.shape(), &[3, 4, 6]); +/// ``` pub trait Dot { /// 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(lhs: &ArrayRef, rhs: &ArrayRef) -> Array +{ + 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(lhs: &ArrayRef, rhs: &ArrayRef, out: &mut ArrayRef) +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::().unwrap(); + let mut out_2d = out.view_mut().into_dimensionality::().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 Dot> for ArrayRef + where A: LinalgScalar + { + type Output = Array; + + #[track_caller] + fn dot(&self, rhs: &ArrayRef) -> Array + { + 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 Dot> for ArrayRef +where A: LinalgScalar +{ + type Output = Array; + + #[track_caller] + fn dot(&self, rhs: &ArrayRef) -> Array + { + nd_dot(&self.view().into_dyn(), rhs) + } +} + +impl_dots!(IxDyn, Ix2); +impl_dots!(IxDyn, IxDyn); + impl Dot> for ArrayRef where A: LinalgScalar { diff --git a/tests/oper.rs b/tests/oper.rs index a6d7054ba..1da5019fc 100644 --- a/tests/oper.rs +++ b/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(lhs: &Array, rhs: &Array2) -> Array +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 = ArrayBuilder::new((3, 4, 5)).build(); + let rhs: Array2 = 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 = ArrayBuilder::new((6, 4, 5)).build(); + let lhs = base.slice(s![..;2, .., ..]).to_owned(); + let rhs: Array2 = 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 = ArrayBuilder::new((2, 3, 4)).build(); + let rhs: Array2 = 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 = Array3::zeros((3, 4, 5)); + let rhs: Array2 = Array2::zeros((6, 7)); + let _ = lhs.dot(&rhs); +} + +#[test] +fn dot_4d_by_2d() +{ + let lhs: Array4 = ArrayBuilder::new((2, 3, 4, 5)).build(); + let rhs: Array2 = 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 = + 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 = 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 = + 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 = 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 = ArrayD::from_shape_vec(IxDyn(&[3, 4, 5]), (0..60).map(|x| x as f64).collect()).unwrap(); + let rhs: Array2 = ArrayBuilder::new((5, 6)).build(); + + let result = lhs.dot(&rhs); + assert_eq!(result.shape(), &[3, 4, 6]); + + let lhs_fixed: Array3 = lhs.into_dimensionality::().unwrap(); + let expected = lhs_fixed.dot(&rhs); + assert_eq!(result.into_dimensionality::().unwrap(), expected); +} + +#[test] +fn dot_dyn_5d_by_2d() +{ + let lhs: ArrayD = + 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 = 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 = lhs.into_dimensionality::().unwrap(); + let expected: Array5 = lhs_fixed.dot(&rhs); + assert_eq!(result.into_dimensionality::().unwrap(), expected); +} + +#[test] +fn dot_dyn_3d_by_2d_non_contiguous() +{ + // Slice with stride 2 to get a non-contiguous layout. + let base: Array3 = ArrayBuilder::new((6, 4, 5)).build(); + let lhs = base.slice(s![..;2, .., ..]).into_dyn(); + let rhs: Array2 = ArrayBuilder::new((5, 7)).build(); + + let result = lhs.dot(&rhs); + assert_eq!(result.shape(), &[3, 4, 7]); + + let lhs_fixed: ArrayView3 = lhs.into_dimensionality::().unwrap(); + let expected = lhs_fixed.dot(&rhs); + assert_eq!(result.into_dimensionality::().unwrap(), expected); +} + +#[test] +#[should_panic(expected = "are not compatible for nd dot")] +fn dot_dyn_shape_mismatch() +{ + let lhs: ArrayD = ArrayD::zeros(IxDyn(&[3, 4, 5])); + let rhs: Array2 = Array2::zeros((6, 7)); + let _ = lhs.dot(&rhs); +}