einsum

Contents

einsum#

class braintrace.einsum(equation, x, weight, *, weight_fn=None)[source]#

ETP-aware two-operand einsum, linear in the trainable weight.

Computes jnp.einsum(equation, x, weight_fn(weight)) through an ETP primitive whose trace rules are derived mechanically from the equation’s axis classification (see the module docstring): weight axes present in the output broadcast (diagonal), weight axes consumed by x are free trace axes (contracted), and output axes absent from the weight are shared (supported when the hidden state carries them, see Notes).

Parameters:

  • equation (str) – A two-operand explicit einsum equation 'x_spec,w_spec->y_spec' in batched form: the leading letter of x_spec must equal the leading letter of y_spec and must not appear in w_spec. v1 restrictions: lowercase letters only, no ellipsis, no repeated letter within one spec, every output letter present in some input, every weight letter present in x or the output. Spaces are stripped before binding.

  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – The non-trainable operand, of rank len(x_spec).

  • weight (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – The trainable operand, of rank len(w_spec). May be a brainunit.Quantity; the units of x and weight multiply into the result.

  • weight_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – Element-wise transform applied to the weight inside the primitive before the contraction. Its Jacobian is composed automatically in the weight-gradient rule.

Returns:

The contraction result, of rank len(y_spec).

Return type:

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

Raises:

ValueError – If the equation violates the v1 restrictions, or if x / weight rank does not match the equation.

See also

matmul

ETP-aware dense matrix multiplication (rank ≤ 2 inputs).

grouped_matmul

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

Notes

There is no bias parameter: compose with a plain add (a plain-add parameter is deliberately excluded from ETP by the selection principle) or use matmul() / grouped_matmul(), which carry one.

Shared-axis equations (output axes absent from the weight spec, e.g. the t in 'btk,kn->btn') are exact for D-RTRL when the hidden state fed by the output carries the shared axes (e.g. a (B, T, N) hidden state) — proven against the BPTT oracle. If the output’s shared-axis positions instead collapse into a smaller hidden state (e.g. rec.sum(axis=1) into a (B, N) hidden), online learning fails loudly at compile time with a cotangent-shape error: the per-position structure required for the weight gradient cannot be recovered from a hidden-shaped learning signal.

Examples

>>> import jax.numpy as jnp
>>> import braintrace
>>> x = jnp.ones((5, 3))
>>> w = jnp.ones((3, 4))
>>> braintrace.einsum('bk,kn->bn', x, w).shape   # dense matmul
(5, 4)
>>>
>>> xh = jnp.ones((5, 2, 3))
>>> wh = jnp.ones((2, 3, 4))
>>> braintrace.einsum('bhd,hde->bhe', xh, wh).shape  # per-head mixing
(5, 2, 4)
Contents