Skip to main content

matmul

Function matmul 

Source
pub fn matmul(
    a: &[f32],
    b: &[f32],
    bias: Option<&[f32]>,
    m: usize,
    k: usize,
    n: usize,
    batch: usize,
    a_stride: usize,
    b_stride: usize,
    trans_a: bool,
    trans_b: bool,
    alpha: f32,
) -> Vec<f32>
Expand description

Batched matmul matching shaders/matmul.wgsl: C[b] = alpha * A[b] @ B[b] (+ bias broadcast over rows). a_stride/b_stride are per-batch element strides (0 broadcasts). With trans_a, A[b] is stored [k, m]; with trans_b, B[b] is stored [n, k] instead of [k, n].