sparse_matmul

Contents

sparse_matmul#

class braintrace.sparse_matmul(x, weight_data, *, sparse_mat, bias=None)[source]#

ETP-aware sparse matrix multiplication.

Computes \(y = x \mathbin{@} \mathrm{sparse}(w) \; (+ b)\), where only the non-zero entries (weight_data) of the fixed sparse pattern are trainable and participate in eligibility-trace computation. Auto-dispatches batched/unbatched based on x.ndim.

Parameters:

  • x (ArrayLike) – Input array.

  • weight_data (ArrayLike) – Sparse-matrix data, i.e. the non-zero values, shape (nnz,).

  • sparse_mat (object) – Sparse-matrix structure (e.g. a saiunit.sparse matrix object). Must expose with_data (substitute new data into the structure) and yw_to_w_transposed (apply the transposed sparse pattern to a trace).

  • bias (ArrayLike or None, optional) – Bias vector. Default None.

Returns:

Output array.

Return type:

ArrayLike

Contents