matmul

Contents

matmul#

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

ETP-aware matrix multiplication.

Computes \(y = x \mathbin{@} \text{weight\_fn}(w) \; (+ \text{bias\_fn}(b))\). The operation is routed through an ETP primitive so the weight (and optional bias) participates in eligibility-trace computation. Auto-dispatches to etp_mm_p (batched) or etp_mv_p (unbatched) based on x.ndim.

The eligibility trace and gradient are always taken w.r.t. the raw weight/bias before any transform is applied.

Parameters:

  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input array, shape (batch, in_features) or (in_features,). Higher-rank x (x.ndim > 2) is rejected with a ValueError: every ETP trace rule assumes one of these two layouts. For genuine tensor contractions use braintrace.einsum().

  • weight (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Weight matrix, shape (in_features, out_features).

  • bias (Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Bias vector, shape (out_features,). 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, shape-preserving transform applied to weight inside the primitive: the op computes y = x @ weight_fn(weight). The eligibility trace and gradient are taken w.r.t. the raw weight. The transform operates on the unitless mantissa; physical units are split off before and recombined after. Default None (identity). Pass a module-level function, not a fresh lambda, if this is called repeatedly (e.g. once per step): weight_fn is stored as a static eqn.params entry hashed by object identity, so two textually identical lambda objects are cache misses and silently retrace every call.

  • 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 inside the primitive (+ bias_fn(bias)). Ignored when bias is None. The transform operates on the unitless mantissa; physical units are split off before and recombined after. Default None. Same re-tracing caveat as weight_fn: pass a module-level function rather than a fresh lambda.

Returns:

Output array, shape (batch, out_features) or (out_features,).

Return type:

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

Examples

>>> import brainstate
>>> import braintrace
>>>
>>> brainstate.environ.set(precision=64)
>>> x = brainstate.random.randn(16, 3)
>>> w = brainstate.random.randn(3, 4)
>>> y = braintrace.matmul(x, w)
>>> print(y.shape)
(16, 4)
>>>
>>> # Unbatched input with a bias term
>>> x1 = brainstate.random.randn(3)
>>> b = brainstate.random.randn(4)
>>> y1 = braintrace.matmul(x1, w, bias=b)
>>> print(y1.shape)
(4,)
>>>
>>> # Constrain weights to be non-negative via a transform
>>> y2 = braintrace.matmul(x, w, weight_fn=lambda ww: ww ** 2)
>>> print(y2.shape)
(16, 4)
Contents