sparse_matmul

class braintrace.sparse_matmul(x, weight_data, *, sparse_mat, bias=None)

ETP-aware sparse matrix multiplication.

Computes \(y = x \mathbin{@} \mathrm{sparse}(w) \; (+ b)\).

Auto-dispatches batched/unbatched based on x.ndim.

Parameters:

• x – Input array.

• weight_data – Sparse-matrix data (non-zero values).

• sparse_mat – The sparse-matrix structure (e.g. a saiunit.sparse matrix object) — must expose with_data and yw_to_w_transposed.

• bias – Optional bias.

Returns:

Output array.