SparseLinear

class braintrace.nn.SparseLinear

Linear layer with sparse weight matrix.

Supports sparse matrices from brainunit.sparse including CSR, CSC, and COO formats. Only the non-zero entries are stored and updated.

Parameters:

• spar_mat (brainunit.sparse.SparseMatrix) – The sparse weight matrix defining the connectivity structure.

• b_init (Callable, ArrayLike, or None, optional) – Bias initializer. If None, no bias is added.

• in_size (int or tuple of int, optional) – The input size. If not provided, inferred from spar_mat.

• name (str, optional) – Name of the module.

• param_type (type, optional) – Type of parameter state. Default is ParamState.

Variables:

• in_size (tuple) – Input feature size.

• out_size (int) – Output feature size.

• spar_mat (brainunit.sparse.SparseMatrix) – The sparse matrix structure.

• weight (ParamState) – Parameter state containing the sparse ‘weight’ data and optionally ‘bias’.

Examples

>>> import braintrace as braintrace
>>> import brainunit as u
>>> import jax.numpy as jnp
>>>
>>> # Create a sparse linear layer with CSR matrix
>>> indices = jnp.array([[0, 1], [1, 2], [2, 0]])
>>> values = jnp.array([1.0, 2.0, 3.0])
>>> spar_mat = u.sparse.CSR((values, indices[:, 1], indices[:, 0]),
...                          shape=(3, 3))
>>> layer = braintrace.nn.SparseLinear(spar_mat, in_size=(3,))
>>> x = jnp.ones((5, 3))
>>> y = layer(x)
>>> y.shape
(5, 3)

update(x)

Apply the sparse linear transform through the ETP sparse_matmul.

The dense data of the sparse weight is routed through braintrace.sparse_matmul(), so it participates in online-learning trace computation.

Parameters:

x (ArrayLike) – Input array, of shape (..., in_size).

Returns:

ArrayLike – The transformed output, of shape (..., out_size).

SparseLinear.__init__(spar_mat, b_init=None, in_size=None, name=None, param_type=<class 'brainstate.ParamState'>)