braintrace.grouped_matmul

Contents

braintrace.grouped_matmul#

braintrace.grouped_matmul(x, weight, bias=None, *, weight_fn=None, bias_fn=None)#

ETP-aware grouped (block-diagonal) matrix multiplication.

Computes y[..., g, :] = x[..., g, :] @ weight_fn(weight)[g] (+ bias_fn(bias)[g]) — a block-diagonal linear map with G independent in_features → out_features blocks, routed through an ETP primitive so weight (and the optional bias) participate in eligibility-trace computation. Auto-dispatches to etp_gmm_p (x.ndim == 3, batched) or etp_gmv_p (x.ndim == 2, unbatched).

Parameters:

  • x (ArrayLike) – Input of shape (groups, in_features) (unbatched) or (batch, groups, in_features) (batched). Any other rank raises ValueError; fold extra leading axes into the batch axis first.

  • weight (ArrayLike) – Block weights of shape (groups, in_features, out_features).

  • bias (ArrayLike, optional) – Per-block bias of shape (groups, out_features). Its unit must be compatible with x.unit * weight.unit.

  • weight_fn (Callable, optional) – Elementwise transform applied to weight inside the primitive (e.g. a sign or non-negativity constraint).

  • bias_fn (Callable, optional) – Elementwise transform applied to bias inside the primitive.

Returns:

ArrayLike – Output of shape (groups, out_features) or (batch, groups, out_features), carrying the product unit when the inputs are brainunit.Quantity.

Raises:

ValueError – If weight is not rank-3, x is not rank-2/rank-3, or the trailing (groups, in_features) axes of x do not match the leading axes of weight.

See also

matmul

the dense ETP matrix multiplication.

Notes

A grouped map is equivalent to a dense linear map whose weight matrix is block_diag(weight[0], ..., weight[G-1]), but its D-RTRL eligibility trace is (batch, G, K, N, n_state) — a factor-G memory reduction over the dense (batch, G·K, G·N, n_state) trace.

Examples

>>> import jax.numpy as jnp
>>> import brainunit as u
>>> import braintrace
>>> x = jnp.ones((16, 4, 8))          # (batch, groups, in)
>>> w = jnp.ones((4, 8, 8))           # (groups, in, out)
>>> y = braintrace.grouped_matmul(x, w)
>>> print(y.shape)
(16, 4, 8)
>>>
>>> # physical-unit quantities are supported
>>> yq = braintrace.grouped_matmul(x * u.mV, w * u.siemens)
>>> print(yq.shape)
(16, 4, 8)
>>>
>>> # constrain block weights to be non-negative via a transform
>>> y2 = braintrace.grouped_matmul(x, w, weight_fn=lambda ww: ww ** 2)
>>> print(y2.shape)
(16, 4, 8)
Contents