Skip to main content

batched_mat_mul_same_shape

Function batched_mat_mul_same_shape 

Source
pub fn batched_mat_mul_same_shape<T>(
    batch: usize,
    m: usize,
    k: usize,
    n: usize,
    a: &[T],
    b: &[T],
) -> Result<Vec<T>, MatrixMulError>
where T: TensorScalar + Copy,
Expand description

Batched matrix multiplication for column-major matrices with one shared shape.

Computes C[p] = A[p] * B[p] for batch matrices. Each A[p] is an m x k column-major matrix and each B[p] is a k x n column-major matrix. The input buffers store complete matrices consecutively, and the returned buffer stores batch consecutive m x n column-major outputs.

§Errors

Returns an error if the input buffer lengths do not match the declared shapes or if the backend rejects the batched GEMM.

§Examples

use tensor4all_tensorbackend::batched_mat_mul_same_shape;

let a = vec![1.0_f64, 3.0, 2.0, 4.0];
let b = vec![5.0_f64, 7.0, 6.0, 8.0];
let out = batched_mat_mul_same_shape(1, 2, 2, 2, &a, &b).unwrap();
assert_eq!(out, vec![19.0, 43.0, 22.0, 50.0]);