GroupedLinear

class braintrace.nn.GroupedLinear(num_groups, in_features, out_features, w_init=KaimingNormal(mode=fan_in, nonlinearity=relu, unit=1), b_init=ZeroInit(unit=1), name=None, param_type=<class 'brainstate.ParamState'>)

Block-diagonal (grouped) linear layer backed by ETP braintrace.grouped_matmul().

Applies num_groups independent in_features → out_features linear maps — equivalent to one dense layer with a block-diagonal weight matrix, but with a num_groups× smaller D-RTRL eligibility trace.

Parameters:

• num_groups (int) – Number of independent blocks G.

• in_features (int) – Input features per block K.

• out_features (int) – Output features per block N.

• w_init (Array | ndarray | bool | number | bool | int | float | complex | Quantity | Callable) – Weight initializer for the (G, K, N) block weights.

• b_init (Array | ndarray | bool | number | bool | int | float | complex | Quantity | Callable | None) – Bias initializer for the (G, N) per-block bias; None disables the bias.

• name (str | None) – Module name.

• param_type (type) – The ParamState subclass holding the parameter dict.

Notes

The underlying op is rank-guarded to x.ndim ∈ {2, 3}; this layer bridges the documented (..., G, K) contract by folding extra leading axes into a single batch axis and unfolding the output.

Examples

>>> import jax.numpy as jnp
>>> import braintrace
>>> layer = braintrace.nn.GroupedLinear(4, 8, 8)
>>> y = layer(jnp.ones((16, 4, 8)))
>>> print(y.shape)
(16, 4, 8)

update(x)

Apply the grouped linear transform through ETP grouped_matmul.

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input array, of shape (..., num_groups, in_features).

Returns:

The transformed output, of shape (..., num_groups, out_features).

Return type:

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