sparse_matmul

Contents

sparse_matmul#

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

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 (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input array.

  • weight (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Sparse-matrix data, i.e. the non-zero values, shape (nnz,).

  • sparse_mat (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), yw_to_w_transposed (apply the transposed sparse pattern to a trace) and yw_to_w (its non-transposed counterpart). Passing any other object raises TypeError.

  • bias (Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Bias vector. Default None.

  • weight_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – 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[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – 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:

Output array.

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

Raises:

TypeError – If sparse_mat is not a brainevent.DataRepresentation.

Contents