From 26fbaee988605c402f0f197ce356fa48c7999f64 Mon Sep 17 00:00:00 2001 From: at264939-ctrl Date: Sun, 5 Jul 2026 21:00:21 +0300 Subject: [PATCH 1/2] implement dot for higher-dimensional arrays with Ix2 --- src/linalg/impl_linalg.rs | 337 +++++++++++++++++++++++--------------- tests/oper.rs | 290 ++++++++++++++++++++------------ 2 files changed, 392 insertions(+), 235 deletions(-) diff --git a/src/linalg/impl_linalg.rs b/src/linalg/impl_linalg.rs index 81c942bc3..1a30ad075 100644 --- a/src/linalg/impl_linalg.rs +++ b/src/linalg/impl_linalg.rs @@ -45,8 +45,7 @@ const GEMM_BLAS_CUTOFF: usize = 7; #[allow(non_camel_case_types)] type blas_index = c_int; // blas index type -impl ArrayRef -{ +impl ArrayRef { /// Perform dot product or matrix multiplication of arrays `self` and `rhs`. /// /// `Rhs` may be either a one-dimensional or a two-dimensional array. @@ -66,13 +65,15 @@ impl ArrayRef /// layout allows. #[track_caller] pub fn dot(&self, rhs: &Rhs) -> >::Output - where Self: Dot + where + Self: Dot, { Dot::dot(self, rhs) } fn dot_generic(&self, rhs: &ArrayRef) -> A - where A: LinalgScalar + where + A: LinalgScalar, { debug_assert_eq!(self.len(), rhs.len()); assert!(self.len() == rhs.len()); @@ -92,14 +93,16 @@ impl ArrayRef #[cfg(not(feature = "blas"))] fn dot_impl(&self, rhs: &ArrayRef) -> A - where A: LinalgScalar + where + A: LinalgScalar, { self.dot_generic(rhs) } #[cfg(feature = "blas")] fn dot_impl(&self, rhs: &ArrayRef) -> A - where A: LinalgScalar + where + A: LinalgScalar, { // Use only if the vector is large enough to be worth it if self.len() >= DOT_BLAS_CUTOFF { @@ -111,15 +114,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); } } @@ -139,8 +135,7 @@ impl ArrayRef /// which agrees with our pointer for non-negative strides, but /// is at the opposite end for negative strides. #[cfg(feature = "blas")] -unsafe fn blas_1d_params(ptr: *const A, len: usize, stride: isize) -> (*const A, blas_index, blas_index) -{ +unsafe fn blas_1d_params(ptr: *const A, len: usize, stride: isize) -> (*const A, blas_index, blas_index) { // [x x x x] // ^--ptr // stride = -1 @@ -157,8 +152,7 @@ unsafe fn blas_1d_params(ptr: *const A, len: usize, stride: isize) -> (*const /// /// For two-dimensional arrays, the dot method computes the matrix /// multiplication. -pub trait Dot -{ +pub trait Dot { /// The result of the operation. /// /// For two-dimensional arrays: a rectangular array. @@ -183,8 +177,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 +189,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 +201,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) } } @@ -223,7 +214,8 @@ impl_dots!(Ix2, Ix1); impl_dots!(Ix2, Ix2); impl Dot> for ArrayRef -where A: LinalgScalar +where + A: LinalgScalar, { type Output = A; @@ -236,14 +228,14 @@ where A: LinalgScalar /// *Note:* If enabled, uses blas `dot` for elements of `f32, f64` when memory /// layout allows. #[track_caller] - fn dot(&self, rhs: &ArrayRef) -> A - { + fn dot(&self, rhs: &ArrayRef) -> A { self.dot_impl(rhs) } } impl Dot> for ArrayRef -where A: LinalgScalar +where + A: LinalgScalar, { type Output = Array; @@ -257,14 +249,12 @@ where A: LinalgScalar /// /// **Panics** if shapes are incompatible. #[track_caller] - fn dot(&self, rhs: &ArrayRef) -> Array - { + fn dot(&self, rhs: &ArrayRef) -> Array { (*rhs.t()).dot(self) } } -impl ArrayRef -{ +impl ArrayRef { /// Perform matrix multiplication of rectangular arrays `self` and `rhs`. /// /// `Rhs` may be either a one-dimensional or a two-dimensional array. @@ -296,19 +286,20 @@ impl ArrayRef /// ``` #[track_caller] pub fn dot(&self, rhs: &Rhs) -> >::Output - where Self: Dot + where + Self: Dot, { Dot::dot(self, rhs) } } impl Dot> for ArrayRef -where A: LinalgScalar +where + A: LinalgScalar, { type Output = Array2; - fn dot(&self, b: &ArrayRef) -> Array2 - { + fn dot(&self, b: &ArrayRef) -> Array2 { let a = self.view(); let b = b.view(); let ((m, k), (k2, n)) = (a.dim(), b.dim()); @@ -331,27 +322,127 @@ 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); + /// Assumes that `m` and `n` are ≤ `isize::MAX`. #[cold] #[inline(never)] -fn dot_shape_error(m: usize, k: usize, k2: usize, n: usize) -> ! -{ +fn dot_shape_error(m: usize, k: usize, k2: usize, n: usize) -> ! { match m.checked_mul(n) { 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); +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 + ); } /// Perform the matrix multiplication of the rectangular array `self` and @@ -364,13 +455,13 @@ fn general_dot_shape_error(m: usize, k: usize, k2: usize, n: usize, c1: usize, c /// /// **Panics** if shapes are incompatible. impl Dot> for ArrayRef -where A: LinalgScalar +where + A: LinalgScalar, { type Output = Array; #[track_caller] - fn dot(&self, rhs: &ArrayRef) -> Array - { + fn dot(&self, rhs: &ArrayRef) -> Array { let ((m, a), n) = (self.dim(), rhs.dim()); if a != n { dot_shape_error(m, a, n, 1); @@ -386,7 +477,8 @@ where A: LinalgScalar } impl ArrayRef -where D: Dimension +where + D: Dimension, { /// Perform the operation `self += alpha * rhs` efficiently, where /// `alpha` is a scalar and `rhs` is another array. This operation is @@ -412,7 +504,8 @@ use self::mat_mul_general as mat_mul_impl; #[cfg(feature = "blas")] fn mat_mul_impl(alpha: A, a: &ArrayRef2, b: &ArrayRef2, beta: A, c: &mut ArrayRef2) -where A: LinalgScalar +where + A: LinalgScalar, { let ((m, k), (k2, n)) = (a.dim(), b.dim()); debug_assert_eq!(k, k2); @@ -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; @@ -498,7 +591,8 @@ where A: LinalgScalar /// C ← α A B + β C fn mat_mul_general(alpha: A, lhs: &ArrayRef2, rhs: &ArrayRef2, beta: A, c: &mut ArrayRef2) -where A: LinalgScalar +where + A: LinalgScalar, { let ((m, k), (_, n)) = (lhs.dim(), rhs.dim()); @@ -631,7 +725,8 @@ where A: LinalgScalar /// `f32, f64` for all memory layouts. #[track_caller] pub fn general_mat_mul(alpha: A, a: &ArrayRef2, b: &ArrayRef2, beta: A, c: &mut ArrayRef2) -where A: LinalgScalar +where + A: LinalgScalar, { let ((m, k), (k2, n)) = (a.dim(), b.dim()); let (m2, n2) = c.dim(); @@ -655,7 +750,8 @@ where A: LinalgScalar #[track_caller] #[allow(clippy::collapsible_if)] pub fn general_mat_vec_mul(alpha: A, a: &ArrayRef2, x: &ArrayRef1, beta: A, y: &mut ArrayRef1) -where A: LinalgScalar +where + A: LinalgScalar, { unsafe { general_mat_vec_mul_impl(alpha, a, x, beta, y.raw_view_mut()) } } @@ -669,9 +765,9 @@ where A: LinalgScalar /// The caller must ensure that the raw view is valid for writing. /// the destination may be uninitialized iff beta is zero. #[allow(clippy::collapsible_else_if)] -unsafe fn general_mat_vec_mul_impl( - alpha: A, a: &ArrayRef2, x: &ArrayRef1, beta: A, y: RawArrayViewMut, -) where A: LinalgScalar +unsafe fn general_mat_vec_mul_impl(alpha: A, a: &ArrayRef2, x: &ArrayRef1, beta: A, y: RawArrayViewMut) +where + A: LinalgScalar, { let ((m, k), k2) = (a.dim(), x.dim()); let m2 = y.dim(); @@ -705,15 +801,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; @@ -747,7 +843,8 @@ unsafe fn general_mat_vec_mul_impl( /// The kronecker product of a LxN matrix A and a MxR matrix B is a (L*M)x(N*R) /// matrix K formed by the block multiplication A_ij * B. pub fn kron(a: &ArrayRef2, b: &ArrayRef2) -> Array -where A: LinalgScalar +where + A: LinalgScalar, { let dimar = a.shape()[0]; let dimac = a.shape()[1]; @@ -773,25 +870,26 @@ where A: LinalgScalar #[inline(always)] /// Return `true` if `A` and `B` are the same type -fn same_type() -> bool -{ +fn same_type() -> bool { TypeId::of::() == TypeId::of::() } // Read pointer to type `A` as type `B`. // // **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::()); +fn cast_as(a: &A) -> B { + 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) } } /// Return the complex in the form of an array [re, im] #[inline] -fn complex_array(z: Complex) -> [A; 2] -{ +fn complex_array(z: Complex) -> [A; 2] { [z.re, z.im] } @@ -817,17 +915,14 @@ where #[cfg(feature = "blas")] #[derive(Copy, Clone)] #[cfg_attr(test, derive(PartialEq, Eq, Debug))] -enum BlasOrder -{ +enum BlasOrder { C, F, } #[cfg(feature = "blas")] -impl BlasOrder -{ - fn transpose(self) -> Self - { +impl BlasOrder { + fn transpose(self) -> Self { match self { Self::C => Self::F, Self::F => Self::C, @@ -836,16 +931,14 @@ impl BlasOrder #[inline] /// Axis of leading stride (opposite of contiguous axis) - fn get_blas_lead_axis(self) -> usize - { + fn get_blas_lead_axis(self) -> usize { match self { Self::C => 0, Self::F => 1, } } - fn to_cblas_layout(self) -> CBLAS_LAYOUT - { + fn to_cblas_layout(self) -> CBLAS_LAYOUT { match self { Self::C => CBLAS_LAYOUT::CblasRowMajor, Self::F => CBLAS_LAYOUT::CblasColMajor, @@ -854,8 +947,7 @@ impl BlasOrder /// When using cblas_sgemm (etc) with C matrix using `for_layout`, /// how should this `self` matrix be transposed - fn to_cblas_transpose_for(self, for_layout: CBLAS_LAYOUT) -> CBLAS_TRANSPOSE - { + fn to_cblas_transpose_for(self, for_layout: CBLAS_LAYOUT) -> CBLAS_TRANSPOSE { let effective_order = match for_layout { CBLAS_LAYOUT::CblasRowMajor => self, CBLAS_LAYOUT::CblasColMajor => self.transpose(), @@ -869,8 +961,7 @@ impl BlasOrder } #[cfg(feature = "blas")] -fn is_blas_2d(dim: &Ix2, stride: &Ix2, order: BlasOrder) -> bool -{ +fn is_blas_2d(dim: &Ix2, stride: &Ix2, order: BlasOrder) -> bool { let (m, n) = dim.into_pattern(); let s0 = stride[0] as isize; let s1 = stride[1] as isize; @@ -907,8 +998,7 @@ fn is_blas_2d(dim: &Ix2, stride: &Ix2, order: BlasOrder) -> bool /// Get BLAS compatible layout if any (C or F, preferring the former) #[cfg(feature = "blas")] -fn get_blas_compatible_layout(a: &ArrayRef) -> Option -{ +fn get_blas_compatible_layout(a: &ArrayRef) -> Option { if is_blas_2d(a._dim(), a._strides(), BlasOrder::C) { Some(BlasOrder::C) } else if is_blas_2d(a._dim(), a._strides(), BlasOrder::F) { @@ -923,8 +1013,7 @@ fn get_blas_compatible_layout(a: &ArrayRef) -> Option /// /// Return leading stride (lda, ldb, ldc) of array #[cfg(feature = "blas")] -fn blas_stride(a: &ArrayRef, order: BlasOrder) -> blas_index -{ +fn blas_stride(a: &ArrayRef, order: BlasOrder) -> blas_index { let axis = order.get_blas_lead_axis(); let other_axis = 1 - axis; let len_this = a.shape()[axis]; @@ -986,12 +1075,12 @@ where /// - The array shapes are incompatible for the operation /// - For vector dot product: the vectors have different lengths impl Dot> for ArrayRef -where A: LinalgScalar +where + A: LinalgScalar, { type Output = Array; - fn dot(&self, rhs: &ArrayRef) -> Self::Output - { + fn dot(&self, rhs: &ArrayRef) -> Self::Output { match (self.ndim(), rhs.ndim()) { (1, 1) => { let a = self.view().into_dimensionality::().unwrap(); @@ -1027,37 +1116,32 @@ where A: LinalgScalar #[cfg(test)] #[cfg(feature = "blas")] -mod blas_tests -{ +mod blas_tests { use super::*; #[test] - fn blas_row_major_2d_normal_matrix() - { + fn blas_row_major_2d_normal_matrix() { let m: Array2 = Array2::zeros((3, 5)); assert!(blas_row_major_2d::(&m)); assert!(!blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_row_matrix() - { + fn blas_row_major_2d_row_matrix() { let m: Array2 = Array2::zeros((1, 5)); assert!(blas_row_major_2d::(&m)); assert!(blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_column_matrix() - { + fn blas_row_major_2d_column_matrix() { let m: Array2 = Array2::zeros((5, 1)); assert!(blas_row_major_2d::(&m)); assert!(blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_transposed_row_matrix() - { + fn blas_row_major_2d_transposed_row_matrix() { let m: Array2 = Array2::zeros((1, 5)); let m_t = m.t(); assert!(blas_row_major_2d::(&m_t)); @@ -1065,8 +1149,7 @@ mod blas_tests } #[test] - fn blas_row_major_2d_transposed_column_matrix() - { + fn blas_row_major_2d_transposed_column_matrix() { let m: Array2 = Array2::zeros((5, 1)); let m_t = m.t(); assert!(blas_row_major_2d::(&m_t)); @@ -1074,16 +1157,14 @@ mod blas_tests } #[test] - fn blas_column_major_2d_normal_matrix() - { + fn blas_column_major_2d_normal_matrix() { let m: Array2 = Array2::zeros((3, 5).f()); assert!(!blas_row_major_2d::(&m)); assert!(blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_skip_rows_ok() - { + fn blas_row_major_2d_skip_rows_ok() { let m: Array2 = Array2::zeros((5, 5)); let mv = m.slice(s![..;2, ..]); assert!(blas_row_major_2d::(&mv)); @@ -1091,8 +1172,7 @@ mod blas_tests } #[test] - fn blas_row_major_2d_skip_columns_fail() - { + fn blas_row_major_2d_skip_columns_fail() { let m: Array2 = Array2::zeros((5, 5)); let mv = m.slice(s![.., ..;2]); assert!(!blas_row_major_2d::(&mv)); @@ -1100,8 +1180,7 @@ mod blas_tests } #[test] - fn blas_col_major_2d_skip_columns_ok() - { + fn blas_col_major_2d_skip_columns_ok() { let m: Array2 = Array2::zeros((5, 5).f()); let mv = m.slice(s![.., ..;2]); assert!(blas_column_major_2d::(&mv)); @@ -1109,8 +1188,7 @@ mod blas_tests } #[test] - fn blas_col_major_2d_skip_rows_fail() - { + fn blas_col_major_2d_skip_rows_fail() { let m: Array2 = Array2::zeros((5, 5).f()); let mv = m.slice(s![..;2, ..]); assert!(!blas_column_major_2d::(&mv)); @@ -1118,8 +1196,7 @@ mod blas_tests } #[test] - fn blas_too_short_stride() - { + fn blas_too_short_stride() { // leading stride must be longer than the other dimension // Example, in a 5 x 5 matrix, the leading stride must be >= 5 for BLAS. diff --git a/tests/oper.rs b/tests/oper.rs index a6d7054ba..48c085c9e 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; @@ -14,8 +16,7 @@ use defmac::defmac; use num_traits::Num; use num_traits::Zero; -fn test_oper(op: &str, a: &[f32], b: &[f32], c: &[f32]) -{ +fn test_oper(op: &str, a: &[f32], b: &[f32], c: &[f32]) { let aa = CowArray::from(arr1(a)); let bb = CowArray::from(arr1(b)); let cc = CowArray::from(arr1(c)); @@ -33,7 +34,8 @@ fn test_oper(op: &str, a: &[f32], b: &[f32], c: &[f32]) } fn test_oper_arr(op: &str, mut aa: CowArray, bb: CowArray, cc: CowArray) -where D: Dimension +where + D: Dimension, { match op { "+" => { @@ -70,8 +72,7 @@ where D: Dimension } #[test] -fn operations() -{ +fn operations() { test_oper("+", &[1.0, 2.0, 3.0, 4.0], &[0.0, 1.0, 2.0, 3.0], &[1.0, 3.0, 5.0, 7.0]); test_oper("-", &[1.0, 2.0, 3.0, 4.0], &[0.0, 1.0, 2.0, 3.0], &[1.0, 1.0, 1.0, 1.0]); test_oper("*", &[1.0, 2.0, 3.0, 4.0], &[0.0, 1.0, 2.0, 3.0], &[0.0, 2.0, 6.0, 12.0]); @@ -81,8 +82,7 @@ fn operations() } #[test] -fn scalar_operations() -{ +fn scalar_operations() { let a = arr0::(1.); let b = rcarr1::(&[1., 1.]); let c = rcarr2(&[[1., 1.], [1., 1.]]); @@ -128,8 +128,7 @@ where } #[test] -fn dot_product() -{ +fn dot_product() { let a = Array::from_iter((0..69).map(|x| x as f32)); let b = &a * 2. - 7.; let dot = 197846.; @@ -165,8 +164,7 @@ fn dot_product() // test that we can dot product with a broadcast array #[test] -fn dot_product_0() -{ +fn dot_product_0() { let a = Array::from_iter((0..69).map(|x| x as f32)); let x = 1.5; let b = aview0(&x); @@ -186,8 +184,7 @@ fn dot_product_0() } #[test] -fn dot_product_neg_stride() -{ +fn dot_product_neg_stride() { // test that we can dot with negative stride let a = Array::from_iter((0..69).map(|x| x as f32)); let b = &a * 2. - 7.; @@ -206,8 +203,7 @@ fn dot_product_neg_stride() } #[test] -fn fold_and_sum() -{ +fn fold_and_sum() { let a = Array::from_iter((0..128).map(|x| x as f32)) .into_shape_with_order((8, 16)) .unwrap(); @@ -248,8 +244,7 @@ fn fold_and_sum() } #[test] -fn product() -{ +fn product() { let step = (2. - 0.5) / 127.; let a = Array::from_iter((0..128).map(|i| 0.5 + step * (i as f64))) .into_shape_with_order((8, 16)) @@ -271,19 +266,16 @@ fn product() } } -fn range_mat(m: Ix, n: Ix) -> Array2 -{ +fn range_mat(m: Ix, n: Ix) -> Array2 { ArrayBuilder::new((m, n)).build() } #[cfg(feature = "approx")] -fn range1_mat64(m: Ix) -> Array1 -{ +fn range1_mat64(m: Ix) -> Array1 { ArrayBuilder::new(m).build() } -fn range_i32(m: Ix, n: Ix) -> Array2 -{ +fn range_i32(m: Ix, n: Ix) -> Array2 { ArrayBuilder::new((m, n)).build() } @@ -318,8 +310,7 @@ where } #[test] -fn mat_mul() -{ +fn mat_mul() { let (m, n, k) = (8, 8, 8); let a = range_mat::(m, n); let b = range_mat::(n, k); @@ -381,8 +372,7 @@ fn mat_mul() // Check that matrix multiplication of contiguous matrices returns a // matrix with the same order #[test] -fn mat_mul_order() -{ +fn mat_mul_order() { let (m, n, k) = (8, 8, 8); let a = range_mat::(m, n); let b = range_mat::(n, k); @@ -401,8 +391,7 @@ fn mat_mul_order() // test matrix multiplication shape mismatch #[test] #[should_panic] -fn mat_mul_shape_mismatch() -{ +fn mat_mul_shape_mismatch() { let (m, k, k2, n) = (8, 8, 9, 8); let a = range_mat::(m, k); let b = range_mat::(k2, n); @@ -412,8 +401,7 @@ fn mat_mul_shape_mismatch() // test matrix multiplication shape mismatch #[test] #[should_panic] -fn mat_mul_shape_mismatch_2() -{ +fn mat_mul_shape_mismatch_2() { let (m, k, k2, n) = (8, 8, 8, 8); let a = range_mat::(m, k); let b = range_mat::(k2, n); @@ -424,8 +412,7 @@ fn mat_mul_shape_mismatch_2() // Check that matrix multiplication // supports broadcast arrays. #[test] -fn mat_mul_broadcast() -{ +fn mat_mul_broadcast() { let (m, n, k) = (16, 16, 16); let a = range_mat::(m, n); let x1 = 1.; @@ -444,8 +431,7 @@ fn mat_mul_broadcast() // Check that matrix multiplication supports reversed axes #[test] -fn mat_mul_rev() -{ +fn mat_mul_rev() { let (m, n, k) = (16, 16, 16); let a = range_mat::(m, n); let b = range_mat::(n, k); @@ -461,8 +447,7 @@ fn mat_mul_rev() // Check that matrix multiplication supports arrays with zero rows or columns #[test] -fn mat_mut_zero_len() -{ +fn mat_mut_zero_len() { defmac!(mat_mul_zero_len range_mat_fn => { for n in 0..4 { for m in 0..4 { @@ -483,8 +468,7 @@ fn mat_mut_zero_len() } #[test] -fn scaled_add() -{ +fn scaled_add() { let a = range_mat(16, 15); let mut b = range_mat(16, 15); b.mapv_inplace(f32::exp); @@ -500,8 +484,7 @@ fn scaled_add() #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] // Very slow on CI/CD machines #[test] -fn scaled_add_2() -{ +fn scaled_add_2() { let beta = -2.3; let sizes = vec![ (4, 4, 1, 4), @@ -539,8 +522,7 @@ fn scaled_add_2() #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] // Very slow on CI/CD machines #[test] -fn scaled_add_3() -{ +fn scaled_add_3() { use approx::assert_relative_eq; use ndarray::{Slice, SliceInfo, SliceInfoElem}; use std::convert::TryFrom; @@ -567,10 +549,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(); @@ -592,8 +571,7 @@ fn scaled_add_3() #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] #[test] -fn gen_mat_mul() -{ +fn gen_mat_mul() { use core::f64; let alpha = -2.3; @@ -637,8 +615,7 @@ fn gen_mat_mul() // Test y = A x where A is f-order #[cfg(feature = "approx")] #[test] -fn gemm_64_1_f() -{ +fn gemm_64_1_f() { let a = range_mat::(64, 64).reversed_axes(); let (m, n) = a.dim(); // m x n times n x 1 == m x 1 @@ -650,8 +627,7 @@ fn gemm_64_1_f() } #[test] -fn gen_mat_mul_i32() -{ +fn gen_mat_mul_i32() { let alpha = -1; let beta = 2; let sizes = if cfg!(miri) { @@ -683,8 +659,7 @@ fn gen_mat_mul_i32() #[cfg(feature = "approx")] #[test] #[cfg_attr(miri, ignore)] // Takes too long -fn gen_mat_vec_mul() -{ +fn gen_mat_vec_mul() { use core::f64; use approx::assert_relative_eq; @@ -710,17 +685,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] { @@ -752,8 +717,7 @@ fn gen_mat_vec_mul() #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] // Very slow on CI/CD machines #[test] -fn vec_mat_mul() -{ +fn vec_mat_mul() { use approx::assert_relative_eq; // simple, slow, correct (hopefully) mat mul @@ -774,17 +738,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] { @@ -813,52 +767,33 @@ fn vec_mat_mul() } #[test] -fn kron_square_f64() -{ +fn kron_square_f64() { let a = arr2(&[[1.0, 0.0], [0.0, 1.0]]); let b = arr2(&[[0.0, 1.0], [1.0, 0.0]]); 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]]), ) } #[test] -fn kron_square_i64() -{ +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] -fn kron_i64() -{ +fn kron_i64() { let a = arr2(&[[1, 0]]); let b = arr2(&[[0, 1], [1, 0]]); let r = arr2(&[[0, 1, 0, 0], [1, 0, 0, 0]]); @@ -869,3 +804,148 @@ 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); +} From e2522870c644aa8ffc13f3eb2798b1d605d078e8 Mon Sep 17 00:00:00 2001 From: at264939-ctrl Date: Sun, 5 Jul 2026 21:18:04 +0300 Subject: [PATCH 2/2] fix(linalg): resolve dyn.rs trait resolution error and nightly formatting --- src/linalg/impl_linalg.rs | 166 +++++++++++++++++++++----------------- tests/oper.rs | 120 +++++++++++++++++---------- 2 files changed, 172 insertions(+), 114 deletions(-) diff --git a/src/linalg/impl_linalg.rs b/src/linalg/impl_linalg.rs index 1a30ad075..c0bd18777 100644 --- a/src/linalg/impl_linalg.rs +++ b/src/linalg/impl_linalg.rs @@ -45,7 +45,8 @@ const GEMM_BLAS_CUTOFF: usize = 7; #[allow(non_camel_case_types)] type blas_index = c_int; // blas index type -impl ArrayRef { +impl ArrayRef +{ /// Perform dot product or matrix multiplication of arrays `self` and `rhs`. /// /// `Rhs` may be either a one-dimensional or a two-dimensional array. @@ -65,15 +66,13 @@ impl ArrayRef { /// layout allows. #[track_caller] pub fn dot(&self, rhs: &Rhs) -> >::Output - where - Self: Dot, + where Self: Dot { Dot::dot(self, rhs) } fn dot_generic(&self, rhs: &ArrayRef) -> A - where - A: LinalgScalar, + where A: LinalgScalar { debug_assert_eq!(self.len(), rhs.len()); assert!(self.len() == rhs.len()); @@ -93,16 +92,14 @@ impl ArrayRef { #[cfg(not(feature = "blas"))] fn dot_impl(&self, rhs: &ArrayRef) -> A - where - A: LinalgScalar, + where A: LinalgScalar { self.dot_generic(rhs) } #[cfg(feature = "blas")] fn dot_impl(&self, rhs: &ArrayRef) -> A - where - A: LinalgScalar, + where A: LinalgScalar { // Use only if the vector is large enough to be worth it if self.len() >= DOT_BLAS_CUTOFF { @@ -135,7 +132,8 @@ impl ArrayRef { /// which agrees with our pointer for non-negative strides, but /// is at the opposite end for negative strides. #[cfg(feature = "blas")] -unsafe fn blas_1d_params(ptr: *const A, len: usize, stride: isize) -> (*const A, blas_index, blas_index) { +unsafe fn blas_1d_params(ptr: *const A, len: usize, stride: isize) -> (*const A, blas_index, blas_index) +{ // [x x x x] // ^--ptr // stride = -1 @@ -152,7 +150,8 @@ unsafe fn blas_1d_params(ptr: *const A, len: usize, stride: isize) -> (*const /// /// For two-dimensional arrays, the dot method computes the matrix /// multiplication. -pub trait Dot { +pub trait Dot +{ /// The result of the operation. /// /// For two-dimensional arrays: a rectangular array. @@ -214,8 +213,7 @@ impl_dots!(Ix2, Ix1); impl_dots!(Ix2, Ix2); impl Dot> for ArrayRef -where - A: LinalgScalar, +where A: LinalgScalar { type Output = A; @@ -228,14 +226,14 @@ where /// *Note:* If enabled, uses blas `dot` for elements of `f32, f64` when memory /// layout allows. #[track_caller] - fn dot(&self, rhs: &ArrayRef) -> A { + fn dot(&self, rhs: &ArrayRef) -> A + { self.dot_impl(rhs) } } impl Dot> for ArrayRef -where - A: LinalgScalar, +where A: LinalgScalar { type Output = Array; @@ -249,12 +247,14 @@ where /// /// **Panics** if shapes are incompatible. #[track_caller] - fn dot(&self, rhs: &ArrayRef) -> Array { + fn dot(&self, rhs: &ArrayRef) -> Array + { (*rhs.t()).dot(self) } } -impl ArrayRef { +impl ArrayRef +{ /// Perform matrix multiplication of rectangular arrays `self` and `rhs`. /// /// `Rhs` may be either a one-dimensional or a two-dimensional array. @@ -286,20 +286,19 @@ impl ArrayRef { /// ``` #[track_caller] pub fn dot(&self, rhs: &Rhs) -> >::Output - where - Self: Dot, + where Self: Dot { Dot::dot(self, rhs) } } impl Dot> for ArrayRef -where - A: LinalgScalar, +where A: LinalgScalar { type Output = Array2; - fn dot(&self, b: &ArrayRef) -> Array2 { + fn dot(&self, b: &ArrayRef) -> Array2 + { let a = self.view(); let b = b.view(); let ((m, k), (k2, n)) = (a.dim(), b.dim()); @@ -384,8 +383,7 @@ impl_dot_nd_ix2!(Ix5, Ix5); impl_dot_nd_ix2!(Ix6, Ix6); impl Dot> for ArrayRef -where - A: LinalgScalar, +where A: LinalgScalar { type Output = Array; @@ -398,7 +396,8 @@ where /// /// **Panics** if `self.ndim() < 2` or if shapes are incompatible. #[track_caller] - fn dot(&self, rhs: &ArrayRef) -> Array { + 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]; @@ -424,11 +423,13 @@ where } impl_dots!(IxDyn, Ix2); +impl_dots!(IxDyn, IxDyn); /// Assumes that `m` and `n` are ≤ `isize::MAX`. #[cold] #[inline(never)] -fn dot_shape_error(m: usize, k: usize, k2: usize, n: usize) -> ! { +fn dot_shape_error(m: usize, k: usize, k2: usize, n: usize) -> ! +{ match m.checked_mul(n) { Some(len) if len <= isize::MAX as usize => {} _ => panic!("ndarray: shape {} × {} overflows isize", m, n), @@ -438,7 +439,8 @@ fn dot_shape_error(m: usize, k: usize, k2: usize, n: usize) -> ! { #[cold] #[inline(never)] -fn general_dot_shape_error(m: usize, k: usize, k2: usize, n: usize, c1: usize, c2: usize) -> ! { +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 @@ -455,13 +457,13 @@ fn general_dot_shape_error(m: usize, k: usize, k2: usize, n: usize, c1: usize, c /// /// **Panics** if shapes are incompatible. impl Dot> for ArrayRef -where - A: LinalgScalar, +where A: LinalgScalar { type Output = Array; #[track_caller] - fn dot(&self, rhs: &ArrayRef) -> Array { + fn dot(&self, rhs: &ArrayRef) -> Array + { let ((m, a), n) = (self.dim(), rhs.dim()); if a != n { dot_shape_error(m, a, n, 1); @@ -477,8 +479,7 @@ where } impl ArrayRef -where - D: Dimension, +where D: Dimension { /// Perform the operation `self += alpha * rhs` efficiently, where /// `alpha` is a scalar and `rhs` is another array. This operation is @@ -504,8 +505,7 @@ use self::mat_mul_general as mat_mul_impl; #[cfg(feature = "blas")] fn mat_mul_impl(alpha: A, a: &ArrayRef2, b: &ArrayRef2, beta: A, c: &mut ArrayRef2) -where - A: LinalgScalar, +where A: LinalgScalar { let ((m, k), (k2, n)) = (a.dim(), b.dim()); debug_assert_eq!(k, k2); @@ -591,8 +591,7 @@ where /// C ← α A B + β C fn mat_mul_general(alpha: A, lhs: &ArrayRef2, rhs: &ArrayRef2, beta: A, c: &mut ArrayRef2) -where - A: LinalgScalar, +where A: LinalgScalar { let ((m, k), (_, n)) = (lhs.dim(), rhs.dim()); @@ -725,8 +724,7 @@ where /// `f32, f64` for all memory layouts. #[track_caller] pub fn general_mat_mul(alpha: A, a: &ArrayRef2, b: &ArrayRef2, beta: A, c: &mut ArrayRef2) -where - A: LinalgScalar, +where A: LinalgScalar { let ((m, k), (k2, n)) = (a.dim(), b.dim()); let (m2, n2) = c.dim(); @@ -750,8 +748,7 @@ where #[track_caller] #[allow(clippy::collapsible_if)] pub fn general_mat_vec_mul(alpha: A, a: &ArrayRef2, x: &ArrayRef1, beta: A, y: &mut ArrayRef1) -where - A: LinalgScalar, +where A: LinalgScalar { unsafe { general_mat_vec_mul_impl(alpha, a, x, beta, y.raw_view_mut()) } } @@ -765,9 +762,9 @@ where /// The caller must ensure that the raw view is valid for writing. /// the destination may be uninitialized iff beta is zero. #[allow(clippy::collapsible_else_if)] -unsafe fn general_mat_vec_mul_impl(alpha: A, a: &ArrayRef2, x: &ArrayRef1, beta: A, y: RawArrayViewMut) -where - A: LinalgScalar, +unsafe fn general_mat_vec_mul_impl( + alpha: A, a: &ArrayRef2, x: &ArrayRef1, beta: A, y: RawArrayViewMut, +) where A: LinalgScalar { let ((m, k), k2) = (a.dim(), x.dim()); let m2 = y.dim(); @@ -843,8 +840,7 @@ where /// The kronecker product of a LxN matrix A and a MxR matrix B is a (L*M)x(N*R) /// matrix K formed by the block multiplication A_ij * B. pub fn kron(a: &ArrayRef2, b: &ArrayRef2) -> Array -where - A: LinalgScalar, +where A: LinalgScalar { let dimar = a.shape()[0]; let dimac = a.shape()[1]; @@ -870,14 +866,16 @@ where #[inline(always)] /// Return `true` if `A` and `B` are the same type -fn same_type() -> bool { +fn same_type() -> bool +{ TypeId::of::() == TypeId::of::() } // Read pointer to type `A` as type `B`. // // **Panics** if `A` and `B` are not the same type -fn cast_as(a: &A) -> B { +fn cast_as(a: &A) -> B +{ assert!( same_type::(), "expect type {} and {} to match", @@ -889,7 +887,8 @@ fn cast_as(a: &A) -> B { /// Return the complex in the form of an array [re, im] #[inline] -fn complex_array(z: Complex) -> [A; 2] { +fn complex_array(z: Complex) -> [A; 2] +{ [z.re, z.im] } @@ -915,14 +914,17 @@ where #[cfg(feature = "blas")] #[derive(Copy, Clone)] #[cfg_attr(test, derive(PartialEq, Eq, Debug))] -enum BlasOrder { +enum BlasOrder +{ C, F, } #[cfg(feature = "blas")] -impl BlasOrder { - fn transpose(self) -> Self { +impl BlasOrder +{ + fn transpose(self) -> Self + { match self { Self::C => Self::F, Self::F => Self::C, @@ -931,14 +933,16 @@ impl BlasOrder { #[inline] /// Axis of leading stride (opposite of contiguous axis) - fn get_blas_lead_axis(self) -> usize { + fn get_blas_lead_axis(self) -> usize + { match self { Self::C => 0, Self::F => 1, } } - fn to_cblas_layout(self) -> CBLAS_LAYOUT { + fn to_cblas_layout(self) -> CBLAS_LAYOUT + { match self { Self::C => CBLAS_LAYOUT::CblasRowMajor, Self::F => CBLAS_LAYOUT::CblasColMajor, @@ -947,7 +951,8 @@ impl BlasOrder { /// When using cblas_sgemm (etc) with C matrix using `for_layout`, /// how should this `self` matrix be transposed - fn to_cblas_transpose_for(self, for_layout: CBLAS_LAYOUT) -> CBLAS_TRANSPOSE { + fn to_cblas_transpose_for(self, for_layout: CBLAS_LAYOUT) -> CBLAS_TRANSPOSE + { let effective_order = match for_layout { CBLAS_LAYOUT::CblasRowMajor => self, CBLAS_LAYOUT::CblasColMajor => self.transpose(), @@ -961,7 +966,8 @@ impl BlasOrder { } #[cfg(feature = "blas")] -fn is_blas_2d(dim: &Ix2, stride: &Ix2, order: BlasOrder) -> bool { +fn is_blas_2d(dim: &Ix2, stride: &Ix2, order: BlasOrder) -> bool +{ let (m, n) = dim.into_pattern(); let s0 = stride[0] as isize; let s1 = stride[1] as isize; @@ -998,7 +1004,8 @@ fn is_blas_2d(dim: &Ix2, stride: &Ix2, order: BlasOrder) -> bool { /// Get BLAS compatible layout if any (C or F, preferring the former) #[cfg(feature = "blas")] -fn get_blas_compatible_layout(a: &ArrayRef) -> Option { +fn get_blas_compatible_layout(a: &ArrayRef) -> Option +{ if is_blas_2d(a._dim(), a._strides(), BlasOrder::C) { Some(BlasOrder::C) } else if is_blas_2d(a._dim(), a._strides(), BlasOrder::F) { @@ -1013,7 +1020,8 @@ fn get_blas_compatible_layout(a: &ArrayRef) -> Option { /// /// Return leading stride (lda, ldb, ldc) of array #[cfg(feature = "blas")] -fn blas_stride(a: &ArrayRef, order: BlasOrder) -> blas_index { +fn blas_stride(a: &ArrayRef, order: BlasOrder) -> blas_index +{ let axis = order.get_blas_lead_axis(); let other_axis = 1 - axis; let len_this = a.shape()[axis]; @@ -1075,12 +1083,12 @@ where /// - The array shapes are incompatible for the operation /// - For vector dot product: the vectors have different lengths impl Dot> for ArrayRef -where - A: LinalgScalar, +where A: LinalgScalar { type Output = Array; - fn dot(&self, rhs: &ArrayRef) -> Self::Output { + fn dot(&self, rhs: &ArrayRef) -> Self::Output + { match (self.ndim(), rhs.ndim()) { (1, 1) => { let a = self.view().into_dimensionality::().unwrap(); @@ -1116,32 +1124,37 @@ where #[cfg(test)] #[cfg(feature = "blas")] -mod blas_tests { +mod blas_tests +{ use super::*; #[test] - fn blas_row_major_2d_normal_matrix() { + fn blas_row_major_2d_normal_matrix() + { let m: Array2 = Array2::zeros((3, 5)); assert!(blas_row_major_2d::(&m)); assert!(!blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_row_matrix() { + fn blas_row_major_2d_row_matrix() + { let m: Array2 = Array2::zeros((1, 5)); assert!(blas_row_major_2d::(&m)); assert!(blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_column_matrix() { + fn blas_row_major_2d_column_matrix() + { let m: Array2 = Array2::zeros((5, 1)); assert!(blas_row_major_2d::(&m)); assert!(blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_transposed_row_matrix() { + fn blas_row_major_2d_transposed_row_matrix() + { let m: Array2 = Array2::zeros((1, 5)); let m_t = m.t(); assert!(blas_row_major_2d::(&m_t)); @@ -1149,7 +1162,8 @@ mod blas_tests { } #[test] - fn blas_row_major_2d_transposed_column_matrix() { + fn blas_row_major_2d_transposed_column_matrix() + { let m: Array2 = Array2::zeros((5, 1)); let m_t = m.t(); assert!(blas_row_major_2d::(&m_t)); @@ -1157,14 +1171,16 @@ mod blas_tests { } #[test] - fn blas_column_major_2d_normal_matrix() { + fn blas_column_major_2d_normal_matrix() + { let m: Array2 = Array2::zeros((3, 5).f()); assert!(!blas_row_major_2d::(&m)); assert!(blas_column_major_2d::(&m)); } #[test] - fn blas_row_major_2d_skip_rows_ok() { + fn blas_row_major_2d_skip_rows_ok() + { let m: Array2 = Array2::zeros((5, 5)); let mv = m.slice(s![..;2, ..]); assert!(blas_row_major_2d::(&mv)); @@ -1172,7 +1188,8 @@ mod blas_tests { } #[test] - fn blas_row_major_2d_skip_columns_fail() { + fn blas_row_major_2d_skip_columns_fail() + { let m: Array2 = Array2::zeros((5, 5)); let mv = m.slice(s![.., ..;2]); assert!(!blas_row_major_2d::(&mv)); @@ -1180,7 +1197,8 @@ mod blas_tests { } #[test] - fn blas_col_major_2d_skip_columns_ok() { + fn blas_col_major_2d_skip_columns_ok() + { let m: Array2 = Array2::zeros((5, 5).f()); let mv = m.slice(s![.., ..;2]); assert!(blas_column_major_2d::(&mv)); @@ -1188,7 +1206,8 @@ mod blas_tests { } #[test] - fn blas_col_major_2d_skip_rows_fail() { + fn blas_col_major_2d_skip_rows_fail() + { let m: Array2 = Array2::zeros((5, 5).f()); let mv = m.slice(s![..;2, ..]); assert!(!blas_column_major_2d::(&mv)); @@ -1196,7 +1215,8 @@ mod blas_tests { } #[test] - fn blas_too_short_stride() { + fn blas_too_short_stride() + { // leading stride must be longer than the other dimension // Example, in a 5 x 5 matrix, the leading stride must be >= 5 for BLAS. diff --git a/tests/oper.rs b/tests/oper.rs index 48c085c9e..b62e7f98b 100644 --- a/tests/oper.rs +++ b/tests/oper.rs @@ -16,7 +16,8 @@ use defmac::defmac; use num_traits::Num; use num_traits::Zero; -fn test_oper(op: &str, a: &[f32], b: &[f32], c: &[f32]) { +fn test_oper(op: &str, a: &[f32], b: &[f32], c: &[f32]) +{ let aa = CowArray::from(arr1(a)); let bb = CowArray::from(arr1(b)); let cc = CowArray::from(arr1(c)); @@ -34,8 +35,7 @@ fn test_oper(op: &str, a: &[f32], b: &[f32], c: &[f32]) { } fn test_oper_arr(op: &str, mut aa: CowArray, bb: CowArray, cc: CowArray) -where - D: Dimension, +where D: Dimension { match op { "+" => { @@ -72,7 +72,8 @@ where } #[test] -fn operations() { +fn operations() +{ test_oper("+", &[1.0, 2.0, 3.0, 4.0], &[0.0, 1.0, 2.0, 3.0], &[1.0, 3.0, 5.0, 7.0]); test_oper("-", &[1.0, 2.0, 3.0, 4.0], &[0.0, 1.0, 2.0, 3.0], &[1.0, 1.0, 1.0, 1.0]); test_oper("*", &[1.0, 2.0, 3.0, 4.0], &[0.0, 1.0, 2.0, 3.0], &[0.0, 2.0, 6.0, 12.0]); @@ -82,7 +83,8 @@ fn operations() { } #[test] -fn scalar_operations() { +fn scalar_operations() +{ let a = arr0::(1.); let b = rcarr1::(&[1., 1.]); let c = rcarr2(&[[1., 1.], [1., 1.]]); @@ -128,7 +130,8 @@ where } #[test] -fn dot_product() { +fn dot_product() +{ let a = Array::from_iter((0..69).map(|x| x as f32)); let b = &a * 2. - 7.; let dot = 197846.; @@ -164,7 +167,8 @@ fn dot_product() { // test that we can dot product with a broadcast array #[test] -fn dot_product_0() { +fn dot_product_0() +{ let a = Array::from_iter((0..69).map(|x| x as f32)); let x = 1.5; let b = aview0(&x); @@ -184,7 +188,8 @@ fn dot_product_0() { } #[test] -fn dot_product_neg_stride() { +fn dot_product_neg_stride() +{ // test that we can dot with negative stride let a = Array::from_iter((0..69).map(|x| x as f32)); let b = &a * 2. - 7.; @@ -203,7 +208,8 @@ fn dot_product_neg_stride() { } #[test] -fn fold_and_sum() { +fn fold_and_sum() +{ let a = Array::from_iter((0..128).map(|x| x as f32)) .into_shape_with_order((8, 16)) .unwrap(); @@ -244,7 +250,8 @@ fn fold_and_sum() { } #[test] -fn product() { +fn product() +{ let step = (2. - 0.5) / 127.; let a = Array::from_iter((0..128).map(|i| 0.5 + step * (i as f64))) .into_shape_with_order((8, 16)) @@ -266,16 +273,19 @@ fn product() { } } -fn range_mat(m: Ix, n: Ix) -> Array2 { +fn range_mat(m: Ix, n: Ix) -> Array2 +{ ArrayBuilder::new((m, n)).build() } #[cfg(feature = "approx")] -fn range1_mat64(m: Ix) -> Array1 { +fn range1_mat64(m: Ix) -> Array1 +{ ArrayBuilder::new(m).build() } -fn range_i32(m: Ix, n: Ix) -> Array2 { +fn range_i32(m: Ix, n: Ix) -> Array2 +{ ArrayBuilder::new((m, n)).build() } @@ -310,7 +320,8 @@ where } #[test] -fn mat_mul() { +fn mat_mul() +{ let (m, n, k) = (8, 8, 8); let a = range_mat::(m, n); let b = range_mat::(n, k); @@ -372,7 +383,8 @@ fn mat_mul() { // Check that matrix multiplication of contiguous matrices returns a // matrix with the same order #[test] -fn mat_mul_order() { +fn mat_mul_order() +{ let (m, n, k) = (8, 8, 8); let a = range_mat::(m, n); let b = range_mat::(n, k); @@ -391,7 +403,8 @@ fn mat_mul_order() { // test matrix multiplication shape mismatch #[test] #[should_panic] -fn mat_mul_shape_mismatch() { +fn mat_mul_shape_mismatch() +{ let (m, k, k2, n) = (8, 8, 9, 8); let a = range_mat::(m, k); let b = range_mat::(k2, n); @@ -401,7 +414,8 @@ fn mat_mul_shape_mismatch() { // test matrix multiplication shape mismatch #[test] #[should_panic] -fn mat_mul_shape_mismatch_2() { +fn mat_mul_shape_mismatch_2() +{ let (m, k, k2, n) = (8, 8, 8, 8); let a = range_mat::(m, k); let b = range_mat::(k2, n); @@ -412,7 +426,8 @@ fn mat_mul_shape_mismatch_2() { // Check that matrix multiplication // supports broadcast arrays. #[test] -fn mat_mul_broadcast() { +fn mat_mul_broadcast() +{ let (m, n, k) = (16, 16, 16); let a = range_mat::(m, n); let x1 = 1.; @@ -431,7 +446,8 @@ fn mat_mul_broadcast() { // Check that matrix multiplication supports reversed axes #[test] -fn mat_mul_rev() { +fn mat_mul_rev() +{ let (m, n, k) = (16, 16, 16); let a = range_mat::(m, n); let b = range_mat::(n, k); @@ -447,7 +463,8 @@ fn mat_mul_rev() { // Check that matrix multiplication supports arrays with zero rows or columns #[test] -fn mat_mut_zero_len() { +fn mat_mut_zero_len() +{ defmac!(mat_mul_zero_len range_mat_fn => { for n in 0..4 { for m in 0..4 { @@ -468,7 +485,8 @@ fn mat_mut_zero_len() { } #[test] -fn scaled_add() { +fn scaled_add() +{ let a = range_mat(16, 15); let mut b = range_mat(16, 15); b.mapv_inplace(f32::exp); @@ -484,7 +502,8 @@ fn scaled_add() { #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] // Very slow on CI/CD machines #[test] -fn scaled_add_2() { +fn scaled_add_2() +{ let beta = -2.3; let sizes = vec![ (4, 4, 1, 4), @@ -522,7 +541,8 @@ fn scaled_add_2() { #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] // Very slow on CI/CD machines #[test] -fn scaled_add_3() { +fn scaled_add_3() +{ use approx::assert_relative_eq; use ndarray::{Slice, SliceInfo, SliceInfoElem}; use std::convert::TryFrom; @@ -571,7 +591,8 @@ fn scaled_add_3() { #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] #[test] -fn gen_mat_mul() { +fn gen_mat_mul() +{ use core::f64; let alpha = -2.3; @@ -615,7 +636,8 @@ fn gen_mat_mul() { // Test y = A x where A is f-order #[cfg(feature = "approx")] #[test] -fn gemm_64_1_f() { +fn gemm_64_1_f() +{ let a = range_mat::(64, 64).reversed_axes(); let (m, n) = a.dim(); // m x n times n x 1 == m x 1 @@ -627,7 +649,8 @@ fn gemm_64_1_f() { } #[test] -fn gen_mat_mul_i32() { +fn gen_mat_mul_i32() +{ let alpha = -1; let beta = 2; let sizes = if cfg!(miri) { @@ -659,7 +682,8 @@ fn gen_mat_mul_i32() { #[cfg(feature = "approx")] #[test] #[cfg_attr(miri, ignore)] // Takes too long -fn gen_mat_vec_mul() { +fn gen_mat_vec_mul() +{ use core::f64; use approx::assert_relative_eq; @@ -717,7 +741,8 @@ fn gen_mat_vec_mul() { #[cfg(feature = "approx")] #[cfg_attr(miri, ignore)] // Very slow on CI/CD machines #[test] -fn vec_mat_mul() { +fn vec_mat_mul() +{ use approx::assert_relative_eq; // simple, slow, correct (hopefully) mat mul @@ -767,7 +792,8 @@ fn vec_mat_mul() { } #[test] -fn kron_square_f64() { +fn kron_square_f64() +{ let a = arr2(&[[1.0, 0.0], [0.0, 1.0]]); let b = arr2(&[[0.0, 1.0], [1.0, 0.0]]); @@ -783,7 +809,8 @@ fn kron_square_f64() { } #[test] -fn kron_square_i64() { +fn kron_square_i64() +{ let a = arr2(&[[1, 0], [0, 1]]); let b = arr2(&[[0, 1], [1, 0]]); @@ -793,7 +820,8 @@ fn kron_square_i64() { } #[test] -fn kron_i64() { +fn kron_i64() +{ let a = arr2(&[[1, 0]]); let b = arr2(&[[0, 1], [1, 0]]); let r = arr2(&[[0, 1, 0, 0], [1, 0, 0, 0]]); @@ -831,7 +859,8 @@ where } #[test] -fn dot_3d_by_2d() { +fn dot_3d_by_2d() +{ let lhs: Array3 = ArrayBuilder::new((3, 4, 5)).build(); let rhs: Array2 = ArrayBuilder::new((5, 6)).build(); @@ -843,7 +872,8 @@ fn dot_3d_by_2d() { } #[test] -fn dot_3d_by_2d_non_contiguous() { +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(); @@ -857,7 +887,8 @@ fn dot_3d_by_2d_non_contiguous() { } #[test] -fn dot_3d_by_2d_integer() { +fn dot_3d_by_2d_integer() +{ let lhs: Array3 = ArrayBuilder::new((2, 3, 4)).build(); let rhs: Array2 = ArrayBuilder::new((4, 5)).build(); @@ -870,14 +901,16 @@ fn dot_3d_by_2d_integer() { #[test] #[should_panic(expected = "not compatible for matrix multiplication")] -fn dot_3d_by_2d_shape_mismatch() { +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() { +fn dot_4d_by_2d() +{ let lhs: Array4 = ArrayBuilder::new((2, 3, 4, 5)).build(); let rhs: Array2 = ArrayBuilder::new((5, 6)).build(); @@ -890,7 +923,8 @@ fn dot_4d_by_2d() { // The shapes here match the NumPy example from issue #1587. #[test] -fn dot_5d_by_2d() { +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(); @@ -903,7 +937,8 @@ fn dot_5d_by_2d() { } #[test] -fn dot_6d_by_2d() { +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(); @@ -916,7 +951,8 @@ fn dot_6d_by_2d() { } #[test] -fn dot_dyn_3d_by_2d() { +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(); @@ -929,7 +965,8 @@ fn dot_dyn_3d_by_2d() { } #[test] -fn dot_dyn_5d_by_2d() { +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(); @@ -944,7 +981,8 @@ fn dot_dyn_5d_by_2d() { #[test] #[should_panic(expected = "not compatible for matrix multiplication")] -fn dot_dyn_shape_mismatch() { +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);