ETPPrimitive

class braintrace.ETPPrimitive(name)

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[..., Any]) – Rule with signature (hidden_dim, trace, **params) -> trace.

Return type:

None

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

Install multiple ETP rules in one call.

Any argument left as None is skipped.

Parameters:

• dt_to_t (Callable[..., Any] | None) – D-RTRL trace propagation rule. Default None.

• xy_to_dw (Callable[..., Any] | None) – Weight-gradient rule. Default None.

• init_drtrl (Callable[..., Any] | None) – D-RTRL trace initialiser. Default None.

• init_pp (Callable[..., Any] | None) – pp_prop (IO-dim) df trace initialiser. Default None.

• fast_path (FastPathRules | None) – 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[..., Any] | None) – 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.

Return type:

None

register_init_drtrl(fn)

Install a D-RTRL trace initialiser.

Parameters:

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

Return type:

None

register_init_pp(fn)

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

Parameters:

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

Return type:

None

register_xy_to_dw(fn)

Install a weight-gradient rule.

Parameters:

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

Return type:

None