braintrace.sparse_matmul

braintrace.sparse_matmul(x, weight, *, sparse_mat, bias=None, weight_fn=None, bias_fn=None)

ETP-aware sparse matrix multiplication.

Computes \(y = x \mathbin{@} \mathrm{sparse}(f(w)) \; (+ g(b))\), where only the non-zero entries (weight) 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 (ArrayLike) – Sparse-matrix data, i.e. the non-zero values, shape (nnz,).

• sparse_mat (brainevent.DataRepresentation) – Sparse-matrix structure (e.g. a brainevent.CSR). Must be a brainevent.DataRepresentation, which implements the ETP online-learning protocol: with_data (substitute new data into the structure), dt2t_transposed (apply the transposed sparse pattern to a trace) and dt2t (its non-transposed counterpart). Passing any other object raises TypeError.

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

• weight_fn (callable or None, optional) – Elementwise transform applied to the non-zero weight data before the matmul. None means identity (no transform). When provided, the transform is applied inside the primitive so that xy_to_dw auto-composes the derivative via jax.vjp, returning the gradient w.r.t. the raw data. The transform operates on the unitless mantissa; physical units are split off before and recombined after.

• bias_fn (callable or None, optional) – Elementwise transform applied to bias before it is added to the output. None means identity. The derivative is composed by the same jax.vjp call as weight_fn. The transform operates on the unitless mantissa; physical units are split off before and recombined after.

Returns:

ArrayLike – Output array.

Raises:

TypeError – If sparse_mat is not a brainevent.DataRepresentation.