{ "cells": [ { "cell_type": "markdown", "id": "a1b2c3d4", "metadata": {}, "source": [ "# Creating Custom ETP Primitives\n" ] }, { "cell_type": "markdown", "id": "b2c3d4e5", "metadata": {}, "source": [ "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](../tutorials/five_primitive_functions.ipynb) for the built-in public operators.\n", "\n", "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](customizing_primitive_transforms.ipynb).\n" ] }, { "cell_type": "code", "execution_count": 1, "id": "c3d4e5f6", "metadata": { "execution": { "iopub.execute_input": "2026-07-09T05:33:29.896943Z", "iopub.status.busy": "2026-07-09T05:33:29.896770Z", "iopub.status.idle": "2026-07-09T05:33:32.318106Z", "shell.execute_reply": "2026-07-09T05:33:32.317189Z" } }, "outputs": [], "source": [ "import brainstate\n", "import jax\n", "import jax.numpy as jnp\n", "\n", "import braintrace" ] }, { "cell_type": "markdown", "id": "e3f4a5b6", "metadata": {}, "source": [ "## Argument Conventions\n", "\n", "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.\n", "\n", "### Invar layout\n", "\n", "| Primitive | `invars[0]` | `invars[1]` | `invars[2]` | `invars[3]` | Static params |\n", "|---|---|---|---|---|---|\n", "| `etp_mm_p` / `etp_mv_p` | input `x` | weight `W` | bias `b` (opt) | — | `has_bias` |\n", "| `etp_elemwise_p` | processed `y` | — | — | — | (none) |\n", "| `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` |\n", "| `etp_sp_mm_p` / `etp_sp_mv_p` | input `x` | weight data | bias `b` (opt) | — | `sparse_mat`, `has_bias` |\n", "| `etp_lora_mm_p` / `etp_lora_mv_p` | input `x` | matrix `B` | matrix `A` | bias `b` (opt) | `alpha`, `has_bias` |\n", "\n", "### `trainable_invars_fn` — the N-trainable-input contract\n", "\n", "Instead of hard-coding a single weight index, each primitive registers a function\n", "\n", "```python\n", "trainable_invars_fn: Callable[[dict], Dict[str, int]]\n", "```\n", "\n", "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.\n", "\n", "Built-in examples:\n", "\n", "| Primitive | `has_bias=False` | `has_bias=True` |\n", "|---|---|---|\n", "| `etp_mm_p` / `etp_mv_p` | `{'weight': 1}` | `{'weight': 1, 'bias': 2}` |\n", "| `etp_conv_p` | `{'weight': 1}` | `{'weight': 1, 'bias': 2}` |\n", "| `etp_sp_mm_p` / `etp_sp_mv_p` | `{'weight': 1}` | `{'weight': 1, 'bias': 2}` |\n", "| `etp_lora_mm_p` / `etp_lora_mv_p` | `{'lora_b': 1, 'lora_a': 2}` | `{'lora_b': 1, 'lora_a': 2, 'bias': 3}` |\n", "| `etp_elemwise_p` | `{'weight': 0}` | — |\n", "\n", "Notes:\n", "- The `has_bias` flag is a static parameter (not a traced value) that controls whether the optional bias argument is present.\n", "- For convolution, all `jax.lax.conv_general_dilated` parameters are passed as static params.\n", "- `x_invar_index` points to the non-trainable input; `etp_elemwise_p` sets it to `None` because the op has no separate input." ] }, { "cell_type": "markdown", "id": "f4a5b6c7", "metadata": {}, "source": [ "## Rule Registries (dict API)\n", "\n", "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.\n", "\n", "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.\n", "\n", "### `ETP_RULES_DT_TO_T` — D-RTRL trace propagation\n", "\n", "```python\n", "dt_to_t(hidden_dim: Array, trace: Dict[str, Array], **static_params) -> Dict[str, Array]\n", "```\n", "\n", "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.\n", "\n", "### `ETP_RULES_XY_TO_DW` — instantaneous hidden-to-weight Jacobian\n", "\n", "```python\n", "xy_to_dw(x: Array, hidden_dim: Array, weights: Dict[str, Array], **static_params) -> Dict[str, Array]\n", "```\n", "\n", "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.\n", "\n", "### `ETP_RULES_INIT_DRTRL` — D-RTRL trace initialiser\n", "\n", "```python\n", "init_drtrl(x_var, y_var, weight_vars: Dict[str, Var], num_hidden_state: int) -> Dict[str, Array]\n", "```\n", "\n", "Returns a zero-filled `Dict[str, Array]` shaped to hold the **parameter-dimension** trace used by `D_RTRL` / `ParamDimVjpAlgorithm`. One leaf per trainable key.\n", "\n", "### `ETP_RULES_INIT_PP` — pp-prop / ES-D-RTRL df-trace initialiser\n", "\n", "```python\n", "init_pp(x_var, y_var, weight_vars: Dict[str, Var], num_hidden_state: int) -> Array\n", "```\n", "\n", "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.\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": 15, "id": "a5b6c7d8", "metadata": { "execution": { "iopub.execute_input": "2026-07-09T05:33:36.761420Z", "iopub.status.busy": "2026-07-09T05:33:36.761268Z", "iopub.status.idle": "2026-07-09T05:33:36.766111Z", "shell.execute_reply": "2026-07-09T05:33:36.765051Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "All ETP primitives:\n", " etp_conv [batched]\n", " etp_einsum [batched]\n", " etp_elemwise\n", " etp_emb [batched]\n", " etp_emb_v\n", " etp_gmm [batched]\n", " etp_gmv\n", " etp_lora_mm [batched]\n", " etp_lora_mv\n", " etp_mm [batched]\n", " etp_mv\n", " etp_sp_mm [batched]\n", " etp_sp_mv\n", "\n", "Trace propagation rules (ETP_RULES_DT_TO_T):\n", " etp_conv\n", " etp_einsum\n", " etp_elemwise\n", " etp_emb\n", " etp_emb_v\n", " etp_gmm\n", " etp_gmv\n", " etp_lora_mm\n", " etp_lora_mv\n", " etp_mm\n", " etp_mv\n", " etp_sp_mm\n", " etp_sp_mv\n", "\n", "Weight gradient rules (ETP_RULES_XY_TO_DW):\n", " etp_conv\n", " etp_einsum\n", " etp_elemwise\n", " etp_emb\n", " etp_emb_v\n", " etp_gmm\n", " etp_gmv\n", " etp_lora_mm\n", " etp_lora_mv\n", " etp_mm\n", " etp_mv\n", " etp_sp_mm\n", " etp_sp_mv\n", "\n", "D-RTRL init rules (ETP_RULES_INIT_DRTRL):\n", " etp_conv\n", " etp_einsum\n", " etp_elemwise\n", " etp_emb\n", " etp_emb_v\n", " etp_gmm\n", " etp_gmv\n", " etp_lora_mm\n", " etp_lora_mv\n", " etp_mm\n", " etp_mv\n", " etp_sp_mm\n", " etp_sp_mv\n", "\n", "pp_prop init rules (ETP_RULES_INIT_PP):\n", " etp_conv\n", " etp_einsum\n", " etp_elemwise\n", " etp_emb\n", " etp_emb_v\n", " etp_gmm\n", " etp_gmv\n", " etp_lora_mm\n", " etp_lora_mv\n", " etp_mm\n", " etp_mv\n", " etp_sp_mm\n", " etp_sp_mv\n" ] } ], "source": [ "from braintrace._op import (\n", " ETP_RULES_DT_TO_T,\n", " ETP_RULES_XY_TO_DW,\n", " ETP_RULES_INIT_DRTRL,\n", " ETP_RULES_INIT_PP,\n", " ETP_PRIMITIVES,\n", " BATCHED_PRIMITIVES,\n", ")\n", "\n", "print(\"All ETP primitives:\")\n", "for p in sorted(ETP_PRIMITIVES, key=lambda p: p.name):\n", " batched_tag = \" [batched]\" if p in BATCHED_PRIMITIVES else \"\"\n", " print(f\" {p.name}{batched_tag}\")\n", "\n", "print(\"\\nTrace propagation rules (ETP_RULES_DT_TO_T):\")\n", "for p in sorted(ETP_RULES_DT_TO_T.keys(), key=lambda p: p.name):\n", " print(f\" {p.name}\")\n", "\n", "print(\"\\nWeight gradient rules (ETP_RULES_XY_TO_DW):\")\n", "for p in sorted(ETP_RULES_XY_TO_DW.keys(), key=lambda p: p.name):\n", " print(f\" {p.name}\")\n", "\n", "print(\"\\nD-RTRL init rules (ETP_RULES_INIT_DRTRL):\")\n", "for p in sorted(ETP_RULES_INIT_DRTRL.keys(), key=lambda p: p.name):\n", " print(f\" {p.name}\")\n", "\n", "print(\"\\npp_prop init rules (ETP_RULES_INIT_PP):\")\n", "for p in sorted(ETP_RULES_INIT_PP.keys(), key=lambda p: p.name):\n", " print(f\" {p.name}\")" ] }, { "cell_type": "markdown", "id": "b6c7d8e9", "metadata": {}, "source": [ "## Creating a Custom Primitive\n", "\n", "Adding a new ETP primitive takes only a few steps. Here we create a **scaled matrix multiplication with an optional bias** as an example:\n", "\n", "$$y = \\text{scale} \\cdot (x \\, @ \\, W) \\; (+ b).$$\n", "\n", "The example exercises the whole dict rule API: both the `weight` and `bias` branches are wired end-to-end." ] }, { "cell_type": "code", "execution_count": 16, "id": "c7d8e9f0", "metadata": { "execution": { "iopub.execute_input": "2026-07-09T05:33:36.768052Z", "iopub.status.busy": "2026-07-09T05:33:36.767886Z", "iopub.status.idle": "2026-07-09T05:33:36.772265Z", "shell.execute_reply": "2026-07-09T05:33:36.771393Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Primitive registered: etp_scaled_mm\n", "Type: ETPPrimitive\n" ] } ], "source": [ "import braintrace\n", "from braintrace import register_primitive\n", "\n", "\n", "# Step 1: Define the implementation.\n", "# Plain JAX function — no special annotations needed.\n", "def _scaled_matmul_impl(*args, scale=1.0, has_bias=False):\n", " x, w = args[0], args[1]\n", " y = scale * (x @ w)\n", " if has_bias:\n", " y = y + args[2]\n", " return y\n", "\n", "\n", "# Step 2: Register as an ETP primitive.\n", "# register_primitive() returns an ``ETPPrimitive`` and auto-derives all\n", "# standard JAX rules (abstract_eval, lowering, JVP, transpose, batching).\n", "# The ``trainable_invars_fn`` / ``x_invar_index`` keywords record the invar\n", "# layout the compiler needs to discover trainable inputs.\n", "def _scaled_trainable_invars(params):\n", " \"\"\"Tell the compiler which invars are trainable.\"\"\"\n", " base = {'weight': 1}\n", " if params.get('has_bias', False):\n", " base['bias'] = 2\n", " return base\n", "\n", "\n", "scaled_mm_p = register_primitive(\n", " 'etp_scaled_mm',\n", " _scaled_matmul_impl,\n", " batched=True,\n", " trainable_invars_fn=_scaled_trainable_invars,\n", " x_invar_index=0,\n", ")\n", "\n", "print(\"Primitive registered:\", scaled_mm_p)\n", "print(\"Type:\", type(scaled_mm_p).__name__)" ] }, { "cell_type": "code", "execution_count": 17, "id": "d8e9f0a1", "metadata": { "execution": { "iopub.execute_input": "2026-07-09T05:33:36.774302Z", "iopub.status.busy": "2026-07-09T05:33:36.774139Z", "iopub.status.idle": "2026-07-09T05:33:36.779757Z", "shell.execute_reply": "2026-07-09T05:33:36.778863Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "dt_to_t registered: True\n", "xy_to_dw registered: True\n", "init_drtrl registered: True\n", "init_pp registered: True\n" ] } ], "source": [ "# Step 3: Register the four ETP-specific rules (dict API).\n", "# Each rule accepts / returns ``Dict[str, Array]`` keyed by the names\n", "# in ``trainable_invars_fn`` — here ``'weight'`` and (optionally) ``'bias'``.\n", "\n", "\n", "def _scaled_dt_to_t(hidden_dim, trace, *, scale=1.0, has_bias=False):\n", " # y = scale * x @ w + b\n", " # -> ∂y/∂w along the \"out\" axis is scaled by `scale`; the y→w chain\n", " # link is still elementwise along `out` axis (singleton at axis=-2).\n", " out = {'weight': trace['weight'] * jnp.expand_dims(hidden_dim, axis=-2) * scale}\n", " if has_bias:\n", " out['bias'] = trace['bias'] * hidden_dim\n", " return out\n", "\n", "\n", "def _scaled_xy_to_dw(x, hidden_dim, weights, *, scale=1.0, has_bias=False):\n", " # Single fused VJP over a dict-valued forward function — returns\n", " # gradients for both 'weight' and 'bias' in one pass.\n", " def _fwd(w_dict):\n", " y = scale * (x @ w_dict['weight'])\n", " if has_bias:\n", " y = y + w_dict['bias']\n", " return y\n", " _, vjp_fn = jax.vjp(_fwd, weights)\n", " return vjp_fn(hidden_dim)[0]\n", "\n", "\n", "def _scaled_init_drtrl(x_var, y_var, weight_vars, num_hidden_state):\n", " \"\"\"D-RTRL parameter-dim trace: one leaf per trainable key.\"\"\"\n", " batch = x_var.aval.shape[0]\n", " out = {\n", " 'weight': jnp.zeros(\n", " (batch, *weight_vars['weight'].aval.shape, num_hidden_state)\n", " )\n", " }\n", " if 'bias' in weight_vars:\n", " out['bias'] = jnp.zeros(\n", " (batch, *weight_vars['bias'].aval.shape, num_hidden_state)\n", " )\n", " return out\n", "\n", "\n", "def _scaled_init_pp(x_var, y_var, weight_vars, num_hidden_state):\n", " \"\"\"pp-prop df trace: single array shaped like the output.\"\"\"\n", " return jnp.zeros(\n", " (*y_var.aval.shape, num_hidden_state),\n", " dtype=y_var.aval.dtype,\n", " )\n", "\n", "\n", "scaled_mm_p.register_etp_rules(\n", " dt_to_t=_scaled_dt_to_t,\n", " xy_to_dw=_scaled_xy_to_dw,\n", " init_drtrl=_scaled_init_drtrl,\n", " init_pp=_scaled_init_pp,\n", ")\n", "\n", "# Each ``register_*`` method also exists as a standalone call; the single\n", "# ``register_etp_rules`` call above installs all four at once.\n", "\n", "print(\"dt_to_t registered: \", scaled_mm_p in ETP_RULES_DT_TO_T)\n", "print(\"xy_to_dw registered: \", scaled_mm_p in ETP_RULES_XY_TO_DW)\n", "print(\"init_drtrl registered:\", scaled_mm_p in ETP_RULES_INIT_DRTRL)\n", "print(\"init_pp registered: \", scaled_mm_p in ETP_RULES_INIT_PP)" ] }, { "cell_type": "code", "execution_count": 18, "id": "e9f0a1b2", "metadata": { "execution": { "iopub.execute_input": "2026-07-09T05:33:36.781431Z", "iopub.status.busy": "2026-07-09T05:33:36.781279Z", "iopub.status.idle": "2026-07-09T05:33:36.924159Z", "shell.execute_reply": "2026-07-09T05:33:36.923260Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Output shape : (4, 5)\n", "Matches 2·xw : True\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "With bias : [6.1 6.1 6.1 6.1 6.1]\n" ] } ], "source": [ "# Step 4: Use the custom primitive via ``primitive.bind()``.\n", "\n", "x = jnp.ones((4, 3))\n", "w = jnp.ones((3, 5))\n", "\n", "y = scaled_mm_p.bind(x, w, scale=2.0, has_bias=False)\n", "y_expected = 2.0 * (x @ w)\n", "\n", "print(\"Output shape :\", y.shape)\n", "print(\"Matches 2·xw :\", bool(jnp.allclose(y, y_expected)))\n", "\n", "# With bias:\n", "b = jnp.full((5,), 0.1)\n", "y_bias = scaled_mm_p.bind(x, w, b, scale=2.0, has_bias=True)\n", "print(\"With bias :\", y_bias[0])" ] }, { "cell_type": "code", "execution_count": 19, "id": "f0a1b2c3", "metadata": { "execution": { "iopub.execute_input": "2026-07-09T05:33:36.925750Z", "iopub.status.busy": "2026-07-09T05:33:36.925603Z", "iopub.status.idle": "2026-07-09T05:33:37.030874Z", "shell.execute_reply": "2026-07-09T05:33:37.029936Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "JIT works: True\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Grad shape: (3, 5)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Vmap output shape: (8, 4, 5)\n" ] } ], "source": [ "# All JAX transformations work automatically for the custom primitive.\n", "\n", "# JIT\n", "y_jit = jax.jit(lambda x, w: scaled_mm_p.bind(x, w, scale=2.0, has_bias=False))(x, w)\n", "print(\"JIT works:\", bool(jnp.allclose(y_jit, y_expected)))\n", "\n", "# Grad\n", "dw = jax.grad(lambda w: jnp.sum(scaled_mm_p.bind(x, w, scale=2.0, has_bias=False)))(w)\n", "print(\"Grad shape:\", dw.shape)\n", "\n", "# Vmap\n", "xs = jnp.ones((8, 4, 3))\n", "ys = jax.vmap(lambda xi: scaled_mm_p.bind(xi, w, scale=2.0, has_bias=False))(xs)\n", "print(\"Vmap output shape:\", ys.shape)" ] }, { "cell_type": "markdown", "id": "a63fbcca", "metadata": {}, "source": [ "> **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." ] }, { "cell_type": "markdown", "id": "af6deb1d", "metadata": {}, "source": [ "## The `gradient_enabled` Flag\n", "\n", "`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.\n", "\n", "| `gradient_enabled` | Compiler behaviour | Example |\n", "|---|---|---|\n", "| `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. |\n", "| `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. |\n", "\n", "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)." ] }, { "cell_type": "markdown", "id": "9391990a", "metadata": {}, "source": [ "## Integrating a Primitive with Online Learning\n", "\n", "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).\n", "\n", "**Rule of thumb**\n", "\n", "| Goal | Use |\n", "|---|---|\n", "| Include a parameter in online learning | `braintrace.matmul(x, W)` (or `conv`, `sparse_matmul`, `lora_matmul`, `element_wise`) |\n", "| Exclude a parameter from online learning | regular JAX op: `x @ W`, `lax.conv_general_dilated`, … |\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": null, "id": "5a9cdb2b", "metadata": { "execution": { "iopub.execute_input": "2026-07-09T05:33:37.032752Z", "iopub.status.busy": "2026-07-09T05:33:37.032565Z", "iopub.status.idle": "2026-07-09T05:33:37.529888Z", "shell.execute_reply": "2026-07-09T05:33:37.528217Z" } }, "outputs": [], "source": [ "import brainstate\n", "\n", "\n", "brainstate.random.seed(19)\n", "\n", "\n", "class TinyRNN(brainstate.nn.Module):\n", " def __init__(self, in_dim=4, hid_dim=6):\n", " super().__init__()\n", " self.in_dim = in_dim\n", " self.hid_dim = hid_dim\n", " # Recurrent weight: ETP-enabled (online learning via D-RTRL).\n", " self.W_rec = brainstate.ParamState(\n", " 0.1 * brainstate.random.normal(size=(hid_dim, hid_dim))\n", " )\n", " # Input weight: plain matmul, learned via BPTT instead.\n", " self.W_in = brainstate.ParamState(\n", " 0.1 * brainstate.random.normal(size=(in_dim, hid_dim))\n", " )\n", "\n", " def init_state(self, batch_size=None, **kwargs):\n", " # ``HiddenState`` is what the ETP compiler traces through.\n", " self.h = brainstate.HiddenState(\n", " jnp.zeros((batch_size or 1, self.hid_dim))\n", " )\n", "\n", " def update(self, x):\n", " # W_in is NOT marked -> excluded from ETP.\n", " input_drive = x @ self.W_in.value\n", " # W_rec IS marked -> included in ETP.\n", " rec_drive = braintrace.matmul(self.h.value, self.W_rec.value)\n", " self.h.value = jax.nn.tanh(input_drive + rec_drive)\n", " return self.h.value\n", "\n", "\n", "model = TinyRNN(in_dim=4, hid_dim=6)\n", "\n", "# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.\n", "alg = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros((2, model.in_dim)), batch_size=2)\n", "\n", "print(\"Compiled ETP relations:\", len(alg.graph.hidden_param_op_relations))\n", "for rel in alg.graph.hidden_param_op_relations:\n", " print(\" primitive =\", rel.primitive.name,\n", " \" trainable keys =\", list(rel.trainable_vars.keys()))" ] }, { "cell_type": "markdown", "id": "a1b2c3d5", "metadata": {}, "source": [ "## Summary\n", "\n", "ETP primitives provide a clean, extensible foundation for online learning in recurrent networks:\n", "\n", "- **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).\n", "\n", "- **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.\n", "\n", "- **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.\n", "\n", "- **All JAX transformations (JIT, grad, vmap, JVP) work automatically** — only the four online-learning-specific rules need hand-writing.\n", "\n", "- **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`.\n", "\n", "- **Brainunit quantities** are handled transparently by every user-facing wrapper.\n", "\n", "Where to look for the math:\n", "\n", "| Rule | Algorithm term | Source with derivation |\n", "|---|---|---|\n", "| `xy_to_dw` | $\\operatorname{diag}(\\mathbf{D}_f^t) \\otimes \\mathbf{x}^t$ | docstrings in `braintrace/_op/{dense,conv,elemwise,sparse,lora}.py` |\n", "| `dt_to_t` | $\\mathbf{D}^t \\boldsymbol{\\epsilon}^{t-1}$ ($y \\to W$ link) | same files |\n", "| `init_drtrl` | param-dim trace shape | same files |\n", "| `init_pp` | output-dim df-trace shape | same files |\n", "\n", "Further reading: `advanced/limitations.ipynb` explains the non-parametric-tail invariant and walks through `GRUCell` (3 Linears, only 2 ETP relations).\n", "\n", "Next, [Customizing Parameter Transforms for ETP Operators](customizing_primitive_transforms.ipynb) shows how transform hooks alter parameter semantics without redefining the registration contract." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.11" } }, "nbformat": 4, "nbformat_minor": 5 }