ETPPrimitive

ETPPrimitive#

class braintrace.ETPPrimitive#

A JAX Primitive with ETP rule registration helpers.

Returned by register_primitive(). Supports every standard JAX primitive operation (bind, def_impl, …) and adds five convenience methods for installing ETP-specific rules into the global registries.

See also

register_primitive

Factory that creates and returns an instance.

Examples

>>> import jax.numpy as jnp
>>> import braintrace
>>>
>>> # Register a primitive whose forward delegates to a standard op.
>>> def my_impl(x, w):
...     return x @ w
>>> my_p = braintrace.register_primitive('etp_demo_mm', my_impl, batched=True)
>>> y = my_p.bind(jnp.ones((2, 3)), jnp.ones((3, 4)))
>>> print(y.shape)
(2, 4)
register_dt_to_t(fn)#

Install a D-RTRL trace propagation rule.

Parameters:

fn (Callable) – Rule with signature (hidden_dim, trace, **params) -> trace.

register_etp_rules(*, dt_to_t=None, xy_to_dw=None, init_drtrl=None, init_pp=None, fast_path=None, pp_x_repr=None, snap_anchor=None, snap_adjacency=None)#

Install multiple ETP rules in one call.

Any argument left as None is skipped.

Parameters:

  • dt_to_t (Callable, optional) – D-RTRL trace propagation rule. Default None.

  • xy_to_dw (Callable, optional) – Weight-gradient rule. Default None.

  • init_drtrl (Callable, optional) – D-RTRL trace initialiser. Default None.

  • init_pp (Callable, optional) – pp_prop (IO-dim) df trace initialiser. Default None.

  • fast_path (FastPathRules, optional) – Closed-form param-dim D-RTRL fast-path kernel bundle (instant / recurrent / solve kernels plus the applicable gate). Registered into ETP_FAST_PATH_RULES. Supplied only by primitives with an elementwise dt_to_t rule (mm / mv / elemwise); None leaves the primitive without a fast path. Default None.

  • pp_x_repr (Callable, optional) – IO-dim x-trace representation rule (x, weight_avals) -> x_repr. Registered into ETP_RULES_PP_X_REPR. Supply it when the raw x is not the operand the op is linear in (e.g. etp_emb_p filters the one-hot encoding of its integer indices); None leaves the IO-dim trace filtering the raw x. Default None.

  • snap_anchor (Callable, optional) – SnAp-n anchor declaration eqn_params -> bool. Registered into ETP_RULES_SNAP_ANCHOR. Declares that the primitive’s trace layout keeps, for every slot, one well-defined hidden position the slot’s instantaneous term lands on – the precondition for widening the trailing state axis into a (neighbour, state) axis. None (the default) leaves the primitive unanchored, so recurrence_scope='sparse_n' rejects it loudly.

  • snap_adjacency (Callable, optional) – SnAp-n position-adjacency rule (eqn_params, size) -> pattern. Registered into ETP_RULES_SNAP_ADJACENCY. Supply it only when the primitive’s cross-position coupling is fully determined by static equation parameters; None (the default) makes the position analysis conservative for this primitive. Default None.

register_init_drtrl(fn)#

Install a D-RTRL trace initialiser.

Parameters:

fn (Callable) – Rule with signature (x_var, y_var, weight_var, num_hidden_state) -> zeros.

register_init_pp(fn)#

Install a pp_prop (IO-dim) df trace initialiser.

Parameters:

fn (Callable) – Rule with signature (x_var, y_var, weight_var, num_hidden_state) -> zeros.

register_xy_to_dw(fn)#

Install a weight-gradient rule.

Parameters:

fn (Callable) – Rule with signature (x, hidden_dim, w, **params) -> dw.

ETPPrimitive.__init__(name)#