diff --git a/src/linalg/impl_linalg.rs b/src/linalg/impl_linalg.rs index 81c942bc3..c0bd18777 100644 --- a/src/linalg/impl_linalg.rs +++ b/src/linalg/impl_linalg.rs @@ -111,15 +111,8 @@ impl ArrayRef unsafe { let (lhs_ptr, n, incx) = blas_1d_params(self._ptr().as_ptr(), self.len(), self.strides()[0]); - let (rhs_ptr, _, incy) = - blas_1d_params(rhs._ptr().as_ptr(), rhs.len(), rhs.strides()[0]); - let ret = blas_sys::$func( - n, - lhs_ptr as *const $ty, - incx, - rhs_ptr as *const $ty, - incy, - ); + let (rhs_ptr, _, incy) = blas_1d_params(rhs._ptr().as_ptr(), rhs.len(), rhs.strides()[0]); + let ret = blas_sys::$func(n, lhs_ptr as *const $ty, incx, rhs_ptr as *const $ty, incy); return cast_as::<$ty, A>(&ret); } } @@ -183,8 +176,7 @@ macro_rules! impl_dots { { type Output = as Dot>>::Output; - fn dot(&self, rhs: &ArrayBase) -> Self::Output - { + fn dot(&self, rhs: &ArrayBase) -> Self::Output { Dot::dot(&**self, &**rhs) } } @@ -196,8 +188,7 @@ macro_rules! impl_dots { { type Output = as Dot>>::Output; - fn dot(&self, rhs: &ArrayRef) -> Self::Output - { + fn dot(&self, rhs: &ArrayRef) -> Self::Output { (**self).dot(rhs) } } @@ -209,8 +200,7 @@ macro_rules! impl_dots { { type Output = as Dot>>::Output; - fn dot(&self, rhs: &ArrayBase) -> Self::Output - { + fn dot(&self, rhs: &ArrayBase) -> Self::Output { self.dot(&**rhs) } } @@ -331,6 +321,110 @@ where A: LinalgScalar } } +/// Implement `Dot>` for a fixed higher-dimensional `ArrayRef`. +/// +/// The last axis of `self` must equal the first axis of `rhs`. All other axes +/// of `self` are preserved in the output, with the last axis replaced by the +/// column count of `rhs`. +/// +/// This mirrors NumPy's behaviour for `x @ y` when `x.ndim > 2`. +macro_rules! impl_dot_nd_ix2 { + ($dim:ty, $larger:ty) => { + impl Dot> for ArrayRef + where + A: LinalgScalar, + { + type Output = Array; + + /// Perform matrix multiplication of `self` and matrix `rhs`. + /// + /// The last axis of `self` must have the same length as the + /// number of rows in `rhs`. The output keeps all leading axes of + /// `self` and replaces the last axis with the number of columns + /// in `rhs`. + /// + /// **Panics** if shapes are incompatible. + #[track_caller] + fn dot(&self, rhs: &ArrayRef) -> Array { + let ndim = self.ndim(); + let k = self.shape()[ndim - 1]; + let (k2, n) = rhs.dim(); + if k != k2 { + dot_shape_error(self.len() / k, k, k2, n); + } + let rows = self.len() / k; + + // Flatten all but the last axis, then do a regular 2-D dot. + let lhs_2d = self + .to_shape((rows, k)) + .expect("ndarray: to_shape failed in nd dot"); + let result_2d: Array2 = lhs_2d.dot(rhs); + + let mut out_shape = <$larger>::zeros(ndim); + for i in 0..ndim - 1 { + out_shape[i] = self.shape()[i]; + } + out_shape[ndim - 1] = n; + + result_2d + .to_shape(out_shape) + .expect("ndarray: to_shape failed reshaping nd dot result") + .into_owned() + } + } + + impl_dots!($dim, Ix2); + }; +} + +impl_dot_nd_ix2!(Ix3, Ix3); +impl_dot_nd_ix2!(Ix4, Ix4); +impl_dot_nd_ix2!(Ix5, Ix5); +impl_dot_nd_ix2!(Ix6, Ix6); + +impl Dot> for ArrayRef +where A: LinalgScalar +{ + type Output = Array; + + /// Perform matrix multiplication of `self` and matrix `rhs`. + /// + /// `self` must have at least 2 dimensions. The last axis of `self` must + /// have the same length as the number of rows in `rhs`. The output keeps + /// all leading axes of `self` and replaces the last with the column count + /// of `rhs`. + /// + /// **Panics** if `self.ndim() < 2` or if shapes are incompatible. + #[track_caller] + fn dot(&self, rhs: &ArrayRef) -> Array + { + let ndim = self.ndim(); + assert!(ndim >= 2, "ndarray: dot requires at least a 2-D array on the left, got {}D", ndim); + let k = self.shape()[ndim - 1]; + let (k2, n) = rhs.dim(); + 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: Array2 = 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); + /// Assumes that `m` and `n` are ≤ `isize::MAX`. #[cold] #[inline(never)] @@ -340,18 +434,17 @@ fn dot_shape_error(m: usize, k: usize, k2: usize, n: usize) -> ! Some(len) if len <= isize::MAX as usize => {} _ => panic!("ndarray: shape {} × {} overflows isize", m, n), } - panic!( - "ndarray: inputs {} × {} and {} × {} are not compatible for matrix multiplication", - m, k, k2, n - ); + panic!("ndarray: inputs {} × {} and {} × {} are not compatible for matrix multiplication", m, k, k2, n); } #[cold] #[inline(never)] fn general_dot_shape_error(m: usize, k: usize, k2: usize, n: usize, c1: usize, c2: usize) -> ! { - panic!("ndarray: inputs {} × {}, {} × {}, and output {} × {} are not compatible for matrix multiplication", - m, k, k2, n, c1, c2); + panic!( + "ndarray: inputs {} × {}, {} × {}, and output {} × {} are not compatible for matrix multiplication", + m, k, k2, n, c1, c2 + ); } /// Perform the matrix multiplication of the rectangular array `self` and @@ -467,17 +560,17 @@ where A: LinalgScalar cblas_layout, a_trans, b_trans, - m as blas_index, // m, rows of Op(a) - n as blas_index, // n, cols of Op(b) - k as blas_index, // k, cols of Op(a) - gemm_scalar_cast!($ty, alpha), // alpha - a._ptr().as_ptr() as *const _, // a - lda, // lda - b._ptr().as_ptr() as *const _, // b - ldb, // ldb - gemm_scalar_cast!($ty, beta), // beta - c._ptr().as_ptr() as *mut _, // c - ldc, // ldc + m as blas_index, // m, rows of Op(a) + n as blas_index, // n, cols of Op(b) + k as blas_index, // k, cols of Op(a) + gemm_scalar_cast!($ty, alpha), // alpha + a._ptr().as_ptr() as *const _, // a + lda, // lda + b._ptr().as_ptr() as *const _, // b + ldb, // ldb + gemm_scalar_cast!($ty, beta), // beta + c._ptr().as_ptr() as *mut _, // c + ldc, // ldc ); } return; @@ -705,15 +798,15 @@ unsafe fn general_mat_vec_mul_impl( blas_sys::$gemv( cblas_layout, a_trans, - m as blas_index, // m, rows of Op(a) - k as blas_index, // n, cols of Op(a) - cast_as(&alpha), // alpha + m as blas_index, // m, rows of Op(a) + k as blas_index, // n, cols of Op(a) + cast_as(&alpha), // alpha a._ptr().as_ptr() as *const _, // a - a_stride, // lda - x_ptr as *const _, // x + a_stride, // lda + x_ptr as *const _, // x x_stride, - cast_as(&beta), // beta - y_ptr as *mut _, // y + cast_as(&beta), // beta + y_ptr as *mut _, // y y_stride, ); return; @@ -783,8 +876,12 @@ fn same_type() -> bool // **Panics** if `A` and `B` are not the same type fn cast_as(a: &A) -> B { - assert!(same_type::(), "expect type {} and {} to match", - std::any::type_name::(), std::any::type_name::()); + assert!( + same_type::(), + "expect type {} and {} to match", + std::any::type_name::(), + std::any::type_name::() + ); unsafe { ::std::ptr::read(a as *const _ as *const B) } } diff --git a/tests/oper.rs b/tests/oper.rs index a6d7054ba..b62e7f98b 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; @@ -567,10 +569,7 @@ fn scaled_add_3() let cslice: Vec = if n == 1 { vec![Slice::from(..).step_by(s2).into()] } else { - vec![ - Slice::from(..).step_by(s1).into(), - Slice::from(..).step_by(s2).into(), - ] + vec![Slice::from(..).step_by(s1).into(), Slice::from(..).step_by(s2).into()] }; let c = range_mat::(n, q).into_shape_with_order(cdim).unwrap(); @@ -710,17 +709,7 @@ fn gen_mat_vec_mul() let alpha = -2.3; let beta = f64::consts::PI; - let sizes = vec![ - (4, 4), - (8, 8), - (17, 15), - (4, 17), - (17, 3), - (19, 18), - (16, 17), - (15, 16), - (67, 63), - ]; + let sizes = vec![(4, 4), (8, 8), (17, 15), (4, 17), (17, 3), (19, 18), (16, 17), (15, 16), (67, 63)]; // test different strides for &s1 in &[1, 2, -1, -2] { for &s2 in &[1, 2, -1, -2] { @@ -774,17 +763,7 @@ fn vec_mat_mul() .unwrap() } - let sizes = vec![ - (4, 4), - (8, 8), - (17, 15), - (4, 17), - (17, 3), - (19, 18), - (16, 17), - (15, 16), - (67, 63), - ]; + let sizes = vec![(4, 4), (8, 8), (17, 15), (4, 17), (17, 3), (19, 18), (16, 17), (15, 16), (67, 63)]; // test different strides for &s1 in &[1, 2, -1, -2] { for &s2 in &[1, 2, -1, -2] { @@ -820,22 +799,12 @@ fn kron_square_f64() assert_eq!( kron(&a, &b), - arr2(&[ - [0.0, 1.0, 0.0, 0.0], - [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 1.0], - [0.0, 0.0, 1.0, 0.0] - ]), + arr2(&[[0.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]), ); assert_eq!( kron(&b, &a), - arr2(&[ - [0.0, 0.0, 1.0, 0.0], - [0.0, 0.0, 0.0, 1.0], - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0] - ]), + arr2(&[[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]), ) } @@ -845,15 +814,9 @@ fn kron_square_i64() let a = arr2(&[[1, 0], [0, 1]]); let b = arr2(&[[0, 1], [1, 0]]); - assert_eq!( - kron(&a, &b), - arr2(&[[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]), - ); + assert_eq!(kron(&a, &b), arr2(&[[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]),); - assert_eq!( - kron(&b, &a), - arr2(&[[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]]), - ) + assert_eq!(kron(&b, &a), arr2(&[[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]]),) } #[test] @@ -869,3 +832,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); +}