Creating Custom ETP Primitives

Creating Custom ETP Primitives#

This chapter starts from a new operation and carries it through the complete ETP registration contract: argument layout, trainable inputs, four rule registries, gradient_enabled, and compiler integration. Begin with Operators for Online Learning for the built-in public operators.

The scope here is primitive creation and registration. To change how an existing operator transforms a stored parameter while preserving exact gradients, use Customizing Parameter Transforms for ETP Operators.

import brainstate
import jax
import jax.numpy as jnp

import braintrace

Argument Conventions#

Every ETP primitive follows specific conventions for its input variables (invars) and static parameters. Understanding these conventions is essential when working with the compiler or adding custom primitives.

Invar layout#

Primitive

invars[0]

invars[1]

invars[2]

invars[3]

Static params

etp_mm_p / etp_mv_p

input x

weight W

bias b (opt)

—

has_bias

etp_elemwise_p

processed y

—

—

—

(none)

etp_conv_p

input x

kernel W

bias b (opt)

—

has_bias, strides, padding, lhs_dilation, rhs_dilation, feature_group_count, batch_group_count, dimension_numbers

etp_sp_mm_p / etp_sp_mv_p

input x

weight data

bias b (opt)

—

sparse_mat, has_bias

etp_lora_mm_p / etp_lora_mv_p

input x

matrix B

matrix A

bias b (opt)

alpha, has_bias

trainable_invars_fn — the N-trainable-input contract#

Instead of hard-coding a single weight index, each primitive registers a function

trainable_invars_fn: Callable[[dict], Dict[str, int]]

which maps the equation’s static params onto {trainable_name: invar_index}. The compiler calls it (via get_trainable_invars) at analysis time to discover every trainable input and to route gradients to the owning ParamState pytree leaf.

Built-in examples:

Primitive

has_bias=False

has_bias=True

etp_mm_p / etp_mv_p

{'weight': 1}

{'weight': 1, 'bias': 2}

etp_conv_p

{'weight': 1}

{'weight': 1, 'bias': 2}

etp_sp_mm_p / etp_sp_mv_p

{'weight': 1}

{'weight': 1, 'bias': 2}

etp_lora_mm_p / etp_lora_mv_p

{'lora_b': 1, 'lora_a': 2}

{'lora_b': 1, 'lora_a': 2, 'bias': 3}

etp_elemwise_p

{'weight': 0}

—

Notes:

  • The has_bias flag is a static parameter (not a traced value) that controls whether the optional bias argument is present.

  • For convolution, all jax.lax.conv_general_dilated parameters are passed as static params.

  • x_invar_index points to the non-trainable input; etp_elemwise_p sets it to None because the op has no separate input.

Rule Registries (dict API)#

ETP uses four global dictionaries to store operation-specific rules. These are the only things that need hand-writing — all standard JAX rules are auto-derived from the implementation function.

All four rules operate on Dict[str, Array] (keyed by the names returned by trainable_invars_fn) — except init_pp, which returns a single output-shaped array because pp-prop factorises the trace as \(\boldsymbol{\epsilon}_f \otimes \boldsymbol{\epsilon}_x\) and only needs one df-tensor per primitive output.

ETP_RULES_DT_TO_T — D-RTRL trace propagation#

dt_to_t(hidden_dim: Array, trace: Dict[str, Array], **static_params) -> Dict[str, Array]

Propagates the hidden-state cotangent \(\partial h/\partial y\) through the \(y \to W\) chain factor of the D-RTRL term \(\mathbf{D}^t \boldsymbol{\epsilon}^{t-1}\). Applied per stored trace key.

ETP_RULES_XY_TO_DW — instantaneous hidden-to-weight Jacobian#

xy_to_dw(x: Array, hidden_dim: Array, weights: Dict[str, Array], **static_params) -> Dict[str, Array]

Returns \(\partial h / \partial W\) for every trainable key. This supplies the \(\operatorname{diag}(\mathbf{D}_f^t) \otimes \mathbf{x}^t\) term in D-RTRL and the solve-time pullback in ES-D-RTRL. Typical implementation: a single fused jax.vjp over a dict-valued forward function.

ETP_RULES_INIT_DRTRL — D-RTRL trace initialiser#

init_drtrl(x_var, y_var, weight_vars: Dict[str, Var], num_hidden_state: int) -> Dict[str, Array]

Returns a zero-filled Dict[str, Array] shaped to hold the parameter-dimension trace used by D_RTRL / ParamDimVjpAlgorithm. One leaf per trainable key.

ETP_RULES_INIT_PP — pp-prop / ES-D-RTRL df-trace initialiser#

init_pp(x_var, y_var, weight_vars: Dict[str, Var], num_hidden_state: int) -> Array

Returns a single zero-filled array shaped to hold the output-dimension df trace used by pp_prop (aliases ES_D_RTRL / IODimVjpAlgorithm). The matching \(\boldsymbol{\epsilon}_x\) factor is managed separately by the executor’s x-trace dictionary.

The two INIT_* registries exist because the two algorithm families factorise the trace differently. Both are required for a primitive that should support both algorithms.

from braintrace._op import (
    ETP_RULES_DT_TO_T,
    ETP_RULES_XY_TO_DW,
    ETP_RULES_INIT_DRTRL,
    ETP_RULES_INIT_PP,
    ETP_PRIMITIVES,
    BATCHED_PRIMITIVES,
)

print("All ETP primitives:")
for p in sorted(ETP_PRIMITIVES, key=lambda p: p.name):
    batched_tag = " [batched]" if p in BATCHED_PRIMITIVES else ""
    print(f"  {p.name}{batched_tag}")

print("\nTrace propagation rules (ETP_RULES_DT_TO_T):")
for p in sorted(ETP_RULES_DT_TO_T.keys(), key=lambda p: p.name):
    print(f"  {p.name}")

print("\nWeight gradient rules (ETP_RULES_XY_TO_DW):")
for p in sorted(ETP_RULES_XY_TO_DW.keys(), key=lambda p: p.name):
    print(f"  {p.name}")

print("\nD-RTRL init rules (ETP_RULES_INIT_DRTRL):")
for p in sorted(ETP_RULES_INIT_DRTRL.keys(), key=lambda p: p.name):
    print(f"  {p.name}")

print("\npp_prop init rules (ETP_RULES_INIT_PP):")
for p in sorted(ETP_RULES_INIT_PP.keys(), key=lambda p: p.name):
    print(f"  {p.name}")

All ETP primitives:
  etp_conv [batched]
  etp_einsum [batched]
  etp_elemwise
  etp_emb [batched]
  etp_emb_v
  etp_gmm [batched]
  etp_gmv
  etp_lora_mm [batched]
  etp_lora_mv
  etp_mm [batched]
  etp_mv
  etp_sp_mm [batched]
  etp_sp_mv

Trace propagation rules (ETP_RULES_DT_TO_T):
  etp_conv
  etp_einsum
  etp_elemwise
  etp_emb
  etp_emb_v
  etp_gmm
  etp_gmv
  etp_lora_mm
  etp_lora_mv
  etp_mm
  etp_mv
  etp_sp_mm
  etp_sp_mv

Weight gradient rules (ETP_RULES_XY_TO_DW):
  etp_conv
  etp_einsum
  etp_elemwise
  etp_emb
  etp_emb_v
  etp_gmm
  etp_gmv
  etp_lora_mm
  etp_lora_mv
  etp_mm
  etp_mv
  etp_sp_mm
  etp_sp_mv

D-RTRL init rules (ETP_RULES_INIT_DRTRL):
  etp_conv
  etp_einsum
  etp_elemwise
  etp_emb
  etp_emb_v
  etp_gmm
  etp_gmv
  etp_lora_mm
  etp_lora_mv
  etp_mm
  etp_mv
  etp_sp_mm
  etp_sp_mv

pp_prop init rules (ETP_RULES_INIT_PP):
  etp_conv
  etp_einsum
  etp_elemwise
  etp_emb
  etp_emb_v
  etp_gmm
  etp_gmv
  etp_lora_mm
  etp_lora_mv
  etp_mm
  etp_mv
  etp_sp_mm
  etp_sp_mv

Creating a Custom Primitive#

Adding a new ETP primitive takes only a few steps. Here we create a scaled matrix multiplication with an optional bias as an example:

\[y = \text{scale} \cdot (x \, @ \, W) \; (+ b).\]

The example exercises the whole dict rule API: both the weight and bias branches are wired end-to-end.

import braintrace
from braintrace import register_primitive


# Step 1: Define the implementation.
# Plain JAX function — no special annotations needed.
def _scaled_matmul_impl(*args, scale=1.0, has_bias=False):
    x, w = args[0], args[1]
    y = scale * (x @ w)
    if has_bias:
        y = y + args[2]
    return y


# Step 2: Register as an ETP primitive.
# register_primitive() returns an ``ETPPrimitive`` and auto-derives all
# standard JAX rules (abstract_eval, lowering, JVP, transpose, batching).
# The ``trainable_invars_fn`` / ``x_invar_index`` keywords record the invar
# layout the compiler needs to discover trainable inputs.
def _scaled_trainable_invars(params):
    """Tell the compiler which invars are trainable."""
    base = {'weight': 1}
    if params.get('has_bias', False):
        base['bias'] = 2
    return base


scaled_mm_p = register_primitive(
    'etp_scaled_mm',
    _scaled_matmul_impl,
    batched=True,
    trainable_invars_fn=_scaled_trainable_invars,
    x_invar_index=0,
)

print("Primitive registered:", scaled_mm_p)
print("Type:", type(scaled_mm_p).__name__)

Primitive registered: etp_scaled_mm
Type: ETPPrimitive

# Step 3: Register the four ETP-specific rules (dict API).
# Each rule accepts / returns ``Dict[str, Array]`` keyed by the names
# in ``trainable_invars_fn`` — here ``'weight'`` and (optionally) ``'bias'``.


def _scaled_dt_to_t(hidden_dim, trace, *, scale=1.0, has_bias=False):
    # y = scale * x @ w + b
    #   -> ∂y/∂w along the "out" axis is scaled by `scale`; the y→w chain
    #      link is still elementwise along `out` axis (singleton at axis=-2).
    out = {'weight': trace['weight'] * jnp.expand_dims(hidden_dim, axis=-2) * scale}
    if has_bias:
        out['bias'] = trace['bias'] * hidden_dim
    return out


def _scaled_xy_to_dw(x, hidden_dim, weights, *, scale=1.0, has_bias=False):
    # Single fused VJP over a dict-valued forward function — returns
    # gradients for both 'weight' and 'bias' in one pass.
    def _fwd(w_dict):
        y = scale * (x @ w_dict['weight'])
        if has_bias:
            y = y + w_dict['bias']
        return y
    _, vjp_fn = jax.vjp(_fwd, weights)
    return vjp_fn(hidden_dim)[0]


def _scaled_init_drtrl(x_var, y_var, weight_vars, num_hidden_state):
    """D-RTRL parameter-dim trace: one leaf per trainable key."""
    batch = x_var.aval.shape[0]
    out = {
        'weight': jnp.zeros(
            (batch, *weight_vars['weight'].aval.shape, num_hidden_state)
        )
    }
    if 'bias' in weight_vars:
        out['bias'] = jnp.zeros(
            (batch, *weight_vars['bias'].aval.shape, num_hidden_state)
        )
    return out


def _scaled_init_pp(x_var, y_var, weight_vars, num_hidden_state):
    """pp-prop df trace: single array shaped like the output."""
    return jnp.zeros(
        (*y_var.aval.shape, num_hidden_state),
        dtype=y_var.aval.dtype,
    )


scaled_mm_p.register_etp_rules(
    dt_to_t=_scaled_dt_to_t,
    xy_to_dw=_scaled_xy_to_dw,
    init_drtrl=_scaled_init_drtrl,
    init_pp=_scaled_init_pp,
)

# Each ``register_*`` method also exists as a standalone call; the single
# ``register_etp_rules`` call above installs all four at once.

print("dt_to_t registered:  ", scaled_mm_p in ETP_RULES_DT_TO_T)
print("xy_to_dw registered: ", scaled_mm_p in ETP_RULES_XY_TO_DW)
print("init_drtrl registered:", scaled_mm_p in ETP_RULES_INIT_DRTRL)
print("init_pp registered:  ", scaled_mm_p in ETP_RULES_INIT_PP)

dt_to_t registered:   True
xy_to_dw registered:  True
init_drtrl registered: True
init_pp registered:   True

# Step 4: Use the custom primitive via ``primitive.bind()``.

x = jnp.ones((4, 3))
w = jnp.ones((3, 5))

y = scaled_mm_p.bind(x, w, scale=2.0, has_bias=False)
y_expected = 2.0 * (x @ w)

print("Output shape :", y.shape)
print("Matches 2·xw :", bool(jnp.allclose(y, y_expected)))

# With bias:
b = jnp.full((5,), 0.1)
y_bias = scaled_mm_p.bind(x, w, b, scale=2.0, has_bias=True)
print("With bias    :", y_bias[0])

Output shape : (4, 5)
Matches 2·xw : True

With bias    : [6.1 6.1 6.1 6.1 6.1]

# All JAX transformations work automatically for the custom primitive.

# JIT
y_jit = jax.jit(lambda x, w: scaled_mm_p.bind(x, w, scale=2.0, has_bias=False))(x, w)
print("JIT works:", bool(jnp.allclose(y_jit, y_expected)))

# Grad
dw = jax.grad(lambda w: jnp.sum(scaled_mm_p.bind(x, w, scale=2.0, has_bias=False)))(w)
print("Grad shape:", dw.shape)

# Vmap
xs = jnp.ones((8, 4, 3))
ys = jax.vmap(lambda xi: scaled_mm_p.bind(xi, w, scale=2.0, has_bias=False))(xs)
print("Vmap output shape:", ys.shape)

JIT works: True

Grad shape: (3, 5)

Vmap output shape: (8, 4, 5)

Compiler integration. Because the registration above already declares trainable_invars_fn and x_invar_index, the primitive is ready to be discovered by the ETP compiler (compile_etrace_graph, D_RTRL, ES_D_RTRL) — no extra steps are required. A primitive registered without trainable_invars_fn still works for direct primitive.bind(), JIT, grad, vmap, and JVP, and the compiler falls back to the single-weight {'weight': 1} layout.

The gradient_enabled Flag#

register_primitive() accepts a gradient_enabled keyword (default False). It controls how the compiler treats this primitive when walking from a weight’s output back to a hidden state.

gradient_enabled

Compiler behaviour

Example

False (default)

Treats the primitive as a tail boundary. A preceding ETP weight whose only path to a hidden state passes through this primitive is excluded from ETP, because per-primitive ETP rules cannot express weight-then-weight composition.

All trainable matmul/conv/sparse/LoRA primitives use this.

True

The primitive is identity-like and may sit on the tail of the y -> h walk. Its presence does not exclude an upstream ETP weight.

Only etp_elemwise_p – intended for gating biases, learnable thresholds, etc.

Use gradient_enabled=True only when the primitive’s xy_to_dw rule is itself an identity-like passthrough; mark all genuinely trainable ops with the default. The “weight -> weight -> hidden” exclusion is what makes per-primitive ETP rules sound – see advanced/limitations.ipynb for a worked example with GRUCell (3 Linears, only 2 ETP relations).

Integrating a Primitive with Online Learning#

Marking a weight operation with a braintrace.* primitive is the only thing a model has to do to opt that parameter into online learning. The compiler then walks the jaxpr, finds every ETP primitive, connects it to the downstream hidden states, and builds the eligibility-trace machinery for either D_RTRL (parameter-dim trace) or ES_D_RTRL / pp_prop (IO-dim trace).

Rule of thumb

Goal

Use

Include a parameter in online learning

braintrace.matmul(x, W) (or conv, sparse_matmul, lora_matmul, element_wise)

Exclude a parameter from online learning

regular JAX op: x @ W, lax.conv_general_dilated, …

The short example below wires a vanilla RNN into D_RTRL: only the recurrent weight is marked with braintrace.matmul, so only it receives an eligibility trace. The input weight uses plain @ and is learned by BPTT through the unrolled scan.

import brainstate


brainstate.random.seed(19)


class TinyRNN(brainstate.nn.Module):
    def __init__(self, in_dim=4, hid_dim=6):
        super().__init__()
        self.in_dim = in_dim
        self.hid_dim = hid_dim
        # Recurrent weight: ETP-enabled (online learning via D-RTRL).
        self.W_rec = brainstate.ParamState(
            0.1 * brainstate.random.normal(size=(hid_dim, hid_dim))
        )
        # Input weight: plain matmul, learned via BPTT instead.
        self.W_in = brainstate.ParamState(
            0.1 * brainstate.random.normal(size=(in_dim, hid_dim))
        )

    def init_state(self, batch_size=None, **kwargs):
        # ``HiddenState`` is what the ETP compiler traces through.
        self.h = brainstate.HiddenState(
            jnp.zeros((batch_size or 1, self.hid_dim))
        )

    def update(self, x):
        # W_in is NOT marked -> excluded from ETP.
        input_drive = x @ self.W_in.value
        # W_rec IS marked -> included in ETP.
        rec_drive = braintrace.matmul(self.h.value, self.W_rec.value)
        self.h.value = jax.nn.tanh(input_drive + rec_drive)
        return self.h.value


model = TinyRNN(in_dim=4, hid_dim=6)

# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.
alg = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros((2, model.in_dim)), batch_size=2)

print("Compiled ETP relations:", len(alg.graph.hidden_param_op_relations))
for rel in alg.graph.hidden_param_op_relations:
    print("   primitive =", rel.primitive.name,
          "  trainable keys =", list(rel.trainable_vars.keys()))

Summary#

ETP primitives provide a clean, extensible foundation for online learning in recurrent networks:

  • 8 built-in primitives cover the most common use cases: dense matmul (mm/mv), element-wise ops, convolution, sparse matmul (mm/mv), and LoRA matmul (mm/mv).

  • Dict rule API — every primitive declares its full set of trainable inputs via trainable_invars_fn, and the four ETP rules consume and return Dict[str, Array]. A single primitive can own several ParamState objects (e.g. weight + bias, or \(B + A + b\) in LoRA) and the executor routes gradients to each in one pass.

  • Custom primitives can be added in a few dozen lines: implement the forward function, call register_primitive (declaring trainable_invars_fn so the compiler can discover it), then hand-write the four ETP rules.

  • All JAX transformations (JIT, grad, vmap, JVP) work automatically — only the four online-learning-specific rules need hand-writing.

  • Parameter selection is primitive-based — every brainstate.ParamState is eligible for ETP, and participation depends only on whether a braintrace.* ETP primitive consumed it. Use gradient_enabled=True exclusively for identity-like ops such as etp_elemwise_p.

  • Brainunit quantities are handled transparently by every user-facing wrapper.

Where to look for the math:

Rule

Algorithm term

Source with derivation

xy_to_dw

\(\operatorname{diag}(\mathbf{D}_f^t) \otimes \mathbf{x}^t\)

docstrings in braintrace/_op/{dense,conv,elemwise,sparse,lora}.py

dt_to_t

\(\mathbf{D}^t \boldsymbol{\epsilon}^{t-1}\) (\(y \to W\) link)

same files

init_drtrl

param-dim trace shape

same files

init_pp

output-dim df-trace shape

same files

Further reading: advanced/limitations.ipynb explains the non-parametric-tail invariant and walks through GRUCell (3 Linears, only 2 ETP relations).

Next, Customizing Parameter Transforms for ETP Operators shows how transform hooks alter parameter semantics without redefining the registration contract.