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..33d08c531 100644 --- a/src/linalg/impl_linalg.rs +++ b/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 Dot> for ArrayRef + where A: LinalgScalar + { + type Output = Array; + + #[track_caller] + fn dot(&self, rhs: &ArrayRef) -> Array + { + 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 Dot> for ArrayRef +where A: LinalgScalar +{ + type Output = Array; + + #[track_caller] + fn dot(&self, rhs: &ArrayRef) -> Array + { + 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 Dot> for ArrayRef where A: LinalgScalar { diff --git a/tests/oper.rs b/tests/oper.rs index a6d7054ba..bf462a597 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,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(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 = "not compatible for matrix multiplication")] +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] +#[should_panic(expected = "not compatible for matrix multiplication")] +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); +}