Customizing Primitive Transforms

Customizing Primitive Transforms#

This is the next chapter after ETP Primitives Deep Dive. There we saw how each ETP primitive marks a weight operation and how a custom primitive (scaled_matmul) is registered. Here we add parameter transform hooks — small elementwise functions applied to a weight inside the primitive — and explain the one design rule that keeps online learning exact in their presence. We close with the optional closed-form fast path and its transform-aware gate.

1. Why parameter transforms, and the one rule#

Models often need a reparametrized weight rather than the raw stored parameter:

  • masked / structured weights: weight_fn=lambda w: w * mask

  • standardized weights: subtract mean / divide by std before use

  • sign-constrained weights: weight_fn=lambda w: w ** 2 (non-negative), jnp.abs, jax.nn.softplus

  • squashed weights: weight_fn=jnp.tanh

Rather than rewriting the parameter and losing the raw value (which the optimizer updates), each built-in ETP op takes an optional, shape-preserving *_fn hook. The op computes the forward pass on the transformed weight \(V = f(W)\) while the eligibility trace and gradient stay attached to the raw weight \(W\).

The load-bearing rule. Let \(f\) be the transform and \(V = f(W)\). The transform Jacobian \(f'(W)\) is applied in exactly one place — the primitive’s xy_to_dw rule, via jax.vjp through the impl. Everything else (yw_to_w, init_drtrl, init_pp) is transform-free: those rules operate on \(\partial h / \partial W_{\text{raw}}\) and are correct as written. Because xy_to_dw builds its VJP through the same impl that applies \(f\), the chain rule threads \(f'\) automatically — no per-transform code is needed.

\[ \frac{\partial h}{\partial W_{\text{raw}}} = \underbrace{\frac{\partial h}{\partial V}}_{\text{closed-form rule}} \cdot \underbrace{f'(W_{\text{raw}})}_{\text{threaded by } \texttt{jax.vjp} \text{ in } \texttt{xy\_to\_dw}} \]

Units note: every wrapper splits a brainunit.Quantity into mantissa + unit and applies the transform to the unitless mantissa, recombining the unit afterwards.

import jax
import jax.numpy as jnp
import brainunit as u
import brainstate

import braintrace

brainstate.environ.set(precision=64)

2. Transforms on the built-in ops#

Each ETP op exposes the hooks that make sense for it. The names differ where the op has more than one trainable factor:

Op

Transform hook(s)

braintrace.matmul(x, weight, bias=None, *, weight_fn=None, bias_fn=None)

weight_fn, bias_fn

braintrace.element_wise(weight, *, weight_fn=None)

weight_fn

braintrace.conv(x, kernel, bias=None, *, ..., kernel_fn=None, bias_fn=None)

kernel_fn (note: not weight_fn), bias_fn

braintrace.sparse_matmul(x, weight, *, sparse_mat, bias=None, weight_fn=None, bias_fn=None)

weight_fn, bias_fn

braintrace.lora_matmul(x, B, A, *, alpha=1.0, bias=None, b_fn=None, a_fn=None, bias_fn=None)

per-factor b_fn, a_fn, plus bias_fn

All hooks default to None (identity), so omitting them reproduces the untransformed op bit-for-bit. The forward pass applies the transform:

# matmul with weight_fn + bias_fn: forward uses the transformed weight/bias.
x = brainstate.random.randn(4, 3)
w = brainstate.random.randn(3, 5)
b = brainstate.random.randn(5)

y = braintrace.matmul(x, w, bias=b, weight_fn=jnp.tanh, bias_fn=jnp.abs)
y_ref = x @ jnp.tanh(w) + jnp.abs(b)
print("matmul forward matches x @ tanh(w) + |b| :", bool(jnp.allclose(y, y_ref)))

matmul forward matches x @ tanh(w) + |b| : True

# element_wise: the single weight_fn squashes the parameter in place.
wg = brainstate.random.randn(5)
print("element_wise matches tanh(w):",
      bool(jnp.allclose(braintrace.element_wise(wg, weight_fn=jnp.tanh), jnp.tanh(wg))))

# conv uses kernel_fn (not weight_fn) for the kernel.
xc = brainstate.random.randn(8, 3, 16)        # NCH (JAX default layout)
k = brainstate.random.randn(4, 3, 5)          # OIH
yc = braintrace.conv(xc, k, strides=(1,), padding='SAME', kernel_fn=lambda ww: ww ** 2)
yc_ref = jax.lax.conv_general_dilated(xc, k ** 2, window_strides=(1,), padding='SAME')
print("conv forward matches conv(x, kernel**2):", bool(jnp.allclose(yc, yc_ref)))

# lora has per-factor b_fn / a_fn.
xl = brainstate.random.randn(16, 8)
B = brainstate.random.randn(8, 2)
A = brainstate.random.randn(2, 4)
yl = braintrace.lora_matmul(xl, B, A, alpha=0.5, a_fn=jnp.tanh)
print("lora forward matches alpha * x @ B @ tanh(A):",
      bool(jnp.allclose(yl, 0.5 * (xl @ B @ jnp.tanh(A)))))

element_wise matches tanh(w): True

conv forward matches conv(x, kernel**2): True

lora forward matches alpha * x @ B @ tanh(A): True

# Units act on the mantissa: the transform sees the unitless value, the unit is
# split off before and recombined after.
x_q = jnp.ones((4, 3)) * u.volt
w_q = jnp.ones((3, 5)) * u.siemens
y_q = braintrace.matmul(x_q, w_q, weight_fn=lambda ww: ww ** 2)
print("unit preserved (V * S = A):", u.get_unit(y_q))

unit preserved (V * S = A): A

3. Correctness check: a transformed weight is handled exactly#

A transform is not an approximation. Because xy_to_dw threads \(f'\) via jax.vjp, an exact online algorithm such as D_RTRL must still reproduce the BPTT gradient element-wise. We verify this with the gradient oracle in braintrace._algorithm.oracle:

  • bptt_param_gradients(factory, inputs) — exact backprop-through-time over the sequence sum-of-squares loss.

  • online_param_gradients(factory, inputs, algo_factory=...) — the same total gradient via the online algorithm’s multi-step path.

  • assert_param_gradients_close(online, bptt, atol=...) — element-wise comparison.

We build a tiny RNN cell whose recurrent map routes through braintrace.matmul(..., weight_fn=jnp.tanh) and assert D_RTRL == BPTT.

from braintrace._algorithm.oracle import (
    bptt_param_gradients,
    online_param_gradients,
    assert_param_gradients_close,
)


class TanhWeightRNN(brainstate.nn.Module):
    """A 1-state-group RNN whose weight is squashed by tanh inside the ETP op."""

    def __init__(self, in_dim=2, hid_dim=4):
        super().__init__()
        self.in_dim = in_dim
        self.hid_dim = hid_dim
        # One weight maps [x ; h] -> h, marked with braintrace.matmul.
        self.W = brainstate.ParamState(
            brainstate.random.randn(in_dim + hid_dim, hid_dim) * 0.2
        )

    def init_state(self, batch_size=None, **kwargs):
        size = (self.hid_dim,) if batch_size is None else (batch_size, self.hid_dim)
        self.h = brainstate.HiddenState(jnp.zeros(size))

    def update(self, x):
        xh = jnp.concatenate([x.reshape(1, -1), self.h.value], axis=-1)
        # weight_fn=tanh: forward uses tanh(W); the trace stays attached to raw W.
        self.h.value = jnp.tanh(braintrace.matmul(xh, self.W.value, weight_fn=jnp.tanh))
        return self.h.value


def factory():
    brainstate.random.seed(0)
    return TanhWeightRNN()


brainstate.random.seed(1)
inputs = brainstate.random.randn(6, 2)  # 6 time steps, in_dim=2

bptt = bptt_param_gradients(factory, inputs)
online = online_param_gradients(
    factory, inputs,
    algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step'),
)
assert_param_gradients_close(online, bptt, atol=1e-4)
print("D_RTRL == BPTT with weight_fn=tanh  ->  transform handled exactly")

D_RTRL == BPTT with weight_fn=tanh  ->  transform handled exactly

4. Adding transforms to a custom primitive#

We now extend the scaled_matmul example from the previous tutorial — \(y = \text{scale} \cdot (x\,@\,W) \;(+ b)\) — with transform hooks. The four ETP rules are unchanged from the untransformed version. Only the impl (and the xy_to_dw forward it differentiates) need to apply weight_fn / bias_fn; because xy_to_dw builds its VJP through that same forward, \(f'\) is threaded for free and D_RTRL stays exact.

from braintrace import register_primitive


# Step 1: impl applies the transforms (this is the ONLY place the forward changes).
def _scaled_matmul_impl(*args, scale=1.0, has_bias=False, weight_fn=None, bias_fn=None):
    x, w = args[0], args[1]
    if weight_fn is not None:
        w = weight_fn(w)
    y = scale * (x @ w)
    if has_bias:
        b = args[2]
        if bias_fn is not None:
            b = bias_fn(b)
        y = y + b
    return y


def _scaled_trainable_invars(params):
    base = {'weight': 1}
    if params.get('has_bias', False):
        base['bias'] = 2
    return base


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

# Step 2: the four ETP rules. Note: yw_to_w / init_* are TRANSFORM-FREE.
# Only xy_to_dw differentiates the impl (which applies f), so f' enters here alone.

def _scaled_yw_to_w(hidden_dim, trace, *, scale=1.0, has_bias=False,
                    weight_fn=None, bias_fn=None):
    # y -> W chain link for y = scale * x @ W: scale on the matching out column.
    # No f' here by design.
    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,
                     weight_fn=None, bias_fn=None):
    # Single fused VJP THROUGH the transform -> gradient w.r.t. the RAW weight,
    # with f' threaded automatically by jax.vjp.
    def _fwd(w_dict):
        w = w_dict['weight']
        if weight_fn is not None:
            w = weight_fn(w)
        y = scale * (x @ w)
        if has_bias:
            b = w_dict['bias']
            if bias_fn is not None:
                b = bias_fn(b)
            y = y + b
        return u.get_mantissa(y)

    _, vjp_fn = jax.vjp(_fwd, weights)
    return jax.tree.map(u.get_mantissa, vjp_fn(hidden_dim)[0])


def _scaled_init_drtrl(x_var, y_var, weight_vars, num_hidden_state):
    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):
    return jnp.zeros((*y_var.aval.shape, num_hidden_state), dtype=y_var.aval.dtype)


scaled_mm_p.register_etp_rules(
    yw_to_w=_scaled_yw_to_w,
    xy_to_dw=_scaled_xy_to_dw,
    init_drtrl=_scaled_init_drtrl,
    init_pp=_scaled_init_pp,
)
print("custom primitive registered with transform-aware xy_to_dw")

custom primitive registered with transform-aware xy_to_dw

# Forward + grad work immediately (auto-derived JAX rules thread the transform).
x = jnp.ones((4, 3))
w = jnp.ones((3, 5))
y = scaled_mm_p.bind(x, w, scale=2.0, has_bias=False, weight_fn=jnp.tanh, bias_fn=None)
print("forward matches 2 * (x @ tanh(w)):", bool(jnp.allclose(y, 2.0 * (x @ jnp.tanh(w)))))

dw = jax.grad(
    lambda w_: jnp.sum(
        scaled_mm_p.bind(x, w_, scale=2.0, has_bias=False, weight_fn=jnp.tanh, bias_fn=None)
    )
)(w)
print("grad shape:", dw.shape)

forward matches 2 * (x @ tanh(w)): True

grad shape: (3, 5)

# And the four unchanged ETP rules give an EXACT D_RTRL gradient with the
# transform active, just like the built-in matmul did in Section 3.
class ScaledTanhRNN(brainstate.nn.Module):
    def __init__(self, in_dim=2, hid_dim=4):
        super().__init__()
        self.in_dim = in_dim
        self.hid_dim = hid_dim
        self.W = brainstate.ParamState(
            brainstate.random.randn(in_dim + hid_dim, hid_dim) * 0.2
        )

    def init_state(self, batch_size=None, **kwargs):
        size = (self.hid_dim,) if batch_size is None else (batch_size, self.hid_dim)
        self.h = brainstate.HiddenState(jnp.zeros(size))

    def update(self, x):
        xh = jnp.concatenate([x.reshape(1, -1), self.h.value], axis=-1)
        r = scaled_mm_p.bind(xh, self.W.value, scale=1.5, has_bias=False,
                             weight_fn=jnp.tanh, bias_fn=None)
        self.h.value = jnp.tanh(r)
        return self.h.value


def custom_factory():
    brainstate.random.seed(0)
    return ScaledTanhRNN()


brainstate.random.seed(1)
inp = brainstate.random.randn(6, 2)
bptt_c = bptt_param_gradients(custom_factory, inp)
online_c = online_param_gradients(
    custom_factory, inp,
    algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step'),
)
assert_param_gradients_close(online_c, bptt_c, atol=1e-4)
print("custom scaled_mm: D_RTRL == BPTT with weight_fn=tanh  ->  rules need no change")

custom scaled_mm: D_RTRL == BPTT with weight_fn=tanh  ->  rules need no change

5. (Advanced) The fast path and its transform-aware gate#

For the elementwise-yw_to_w primitives (etp_mm / etp_mv / etp_elemwise), a closed-form param-dim D-RTRL kernel bundle replaces the generic nested-vmap trace path with direct einsums. The bundle is a FastPathRules(instant, recurrent, solve, applicable) namedtuple, registered alongside the four rules via register_etp_rules(fast_path=...), and looked up with get_fast_path_rules(primitive).

  • instant — the instantaneous \(\operatorname{diag}(\mathbf{D}_f^t) \otimes \mathbf{x}^t\) term.

  • recurrent — the \(\mathbf{D}^t \boldsymbol{\epsilon}^{t-1}\) contraction.

  • solve — the solve-time contraction of the learning signal with the trace.

  • applicable(eqn_params) — the gate.

The closed-form kernels compute \(\partial h / \partial V\) for the transformed weight \(V = f(W)\) — they emit the bare outer product and drop \(f'\). So whenever a transform hook is present, applicable returns False, and that relation falls back to the rule path (whose xy_to_dw applies \(f'\) via jax.vjp, as in Sections 3-4). This is what keeps the fast path an optimization, never a correctness hazard.

Note sparse / conv / lora have no fast path at all — they always use the rule path.

from braintrace._op import (
    get_fast_path_rules,
    etp_mm_p, etp_elemwise_p,
    etp_sp_mm_p, etp_conv_p, etp_lora_mm_p,
)

fp = get_fast_path_rules(etp_mm_p)
assert fp is not None
print("etp_mm has a fast path:", fp is not None)

# Gate ON when no transform; OFF the moment a transform hook is present.
print("applicable (no transform)     :", fp.applicable({'weight_fn': None, 'bias_fn': None}))
print("applicable (weight_fn=tanh)    :", fp.applicable({'weight_fn': jnp.tanh, 'bias_fn': None}))
assert fp.applicable({'weight_fn': None, 'bias_fn': None}) is True
assert fp.applicable({'weight_fn': jnp.tanh, 'bias_fn': None}) is False

# elemwise also has one; sparse / conv / lora do not.
print("etp_elemwise fast path:", get_fast_path_rules(etp_elemwise_p) is not None)
for p, name in [(etp_sp_mm_p, 'etp_sp_mm'), (etp_conv_p, 'etp_conv'), (etp_lora_mm_p, 'etp_lora_mm')]:
    print(f"{name} fast path is None:", get_fast_path_rules(p) is None)

etp_mm has a fast path: True
applicable (no transform)     : True
applicable (weight_fn=tanh)    : False
etp_elemwise fast path: True
etp_sp_mm fast path is None: True
etp_conv fast path is None: True
etp_lora_mm fast path is None: True

Summary#

  • Transform hooks let a model use a reparametrized weight (weight_fn / kernel_fn / b_fn / a_fn / bias_fn) while the trace and gradient stay attached to the raw parameter.

  • One rule: the transform Jacobian \(f'\) enters only in xy_to_dw (via jax.vjp through the impl). yw_to_w and the trace initializers stay transform-free and exact.

  • Exactness: D_RTRL == BPTT element-wise even with weight_fn=tanh, for both built-in and custom primitives — verified via braintrace._algorithm.oracle.

  • Custom primitives need no rule changes for transforms: apply the hook in the impl, and xy_to_dw’s VJP threads \(f'\) for free.

  • Fast path: the closed-form FastPathRules bundle drops \(f'\), so its applicable gate disables it under any transform; such relations fall back to the (correct) rule path. sparse / conv / lora have no fast path.