Skip to content

implement dot for higher-dimensional arrays with Ix2 #1602

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
EbrahimEldesoky wants to merge 2 commits into
base: master
Choose a base branch
from
+302 −87
Open

implement dot for higher-dimensional arrays with Ix2 #1602

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
26fbaee
implement dot for higher-dimensional arrays with Ix2
EbrahimEldesoky Jul 5, 2026
e252287
fix(linalg): resolve dyn.rs trait resolution error and nightly format…
EbrahimEldesoky Jul 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 138 additions & 41 deletions src/linalg/impl_linalg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,15 +111,8 @@ impl<A> ArrayRef<A, Ix1>
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);
}
}
Expand Down Expand Up @@ -183,8 +176,7 @@ macro_rules! impl_dots {
{
type Output = <ArrayRef<A, $shape1> as Dot<ArrayRef<A, $shape2>>>::Output;

fn dot(&self, rhs: &ArrayBase<S2, $shape2>) -> Self::Output
{
fn dot(&self, rhs: &ArrayBase<S2, $shape2>) -> Self::Output {
Dot::dot(&**self, &**rhs)
}
}
Expand All @@ -196,8 +188,7 @@ macro_rules! impl_dots {
{
type Output = <ArrayRef<A, $shape1> as Dot<ArrayRef<A, $shape2>>>::Output;

fn dot(&self, rhs: &ArrayRef<A, $shape2>) -> Self::Output
{
fn dot(&self, rhs: &ArrayRef<A, $shape2>) -> Self::Output {
(**self).dot(rhs)
}
}
Expand All @@ -209,8 +200,7 @@ macro_rules! impl_dots {
{
type Output = <ArrayRef<A, $shape1> as Dot<ArrayRef<A, $shape2>>>::Output;

fn dot(&self, rhs: &ArrayBase<S, $shape2>) -> Self::Output
{
fn dot(&self, rhs: &ArrayBase<S, $shape2>) -> Self::Output {
self.dot(&**rhs)
}
}
Expand Down Expand Up @@ -331,6 +321,110 @@ where A: LinalgScalar
}
}

/// Implement `Dot<ArrayRef<A, Ix2>>` for a fixed higher-dimensional `ArrayRef<A, D>`.
///
/// 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<A> Dot<ArrayRef<A, Ix2>> for ArrayRef<A, $dim>
where
A: LinalgScalar,
{
type Output = Array<A, $larger>;

/// 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<A, Ix2>) -> Array<A, $larger> {
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<A> = 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<A> Dot<ArrayRef<A, Ix2>> for ArrayRef<A, IxDyn>
where A: LinalgScalar
{
type Output = Array<A, IxDyn>;

/// 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<A, Ix2>) -> Array<A, IxDyn>
{
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<A> = 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)]
Expand All @@ -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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -705,15 +798,15 @@ unsafe fn general_mat_vec_mul_impl<A>(
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;
Expand Down Expand Up @@ -783,8 +876,12 @@ fn same_type<A: 'static, B: 'static>() -> bool
// **Panics** if `A` and `B` are not the same type
fn cast_as<A: 'static + Copy, B: 'static + Copy>(a: &A) -> B
{
assert!(same_type::<A, B>(), "expect type {} and {} to match",
std::any::type_name::<A>(), std::any::type_name::<B>());
assert!(
same_type::<A, B>(),
"expect type {} and {} to match",
std::any::type_name::<A>(),
std::any::type_name::<B>()
);
unsafe { ::std::ptr::read(a as *const _ as *const B) }
}

Expand Down
Loading
Loading