diff --git a/changelog.md b/changelog.md index 4044ea1..c59d4cf 100644 --- a/changelog.md +++ b/changelog.md @@ -18,6 +18,17 @@ `brainstate._compatible_import.wrap_init` shim (re-exported as `braintrace._compatible_imports.wrap_init`), matching the pattern already used elsewhere in the `brainstate`/`saiunit` stack. No behavior change. +- **Fixed stale API references in the documentation notebooks.** Three + notebooks (`docs/tutorials/etp_primitives.ipynb`, + `docs/tutorials/customizing_primitive_transforms.ipynb`, + `docs/advanced/limitations.ipynb`) still referenced braintrace's own + pre-#130 `yw_to_w` / `ETP_RULES_YW_TO_W` rule naming instead of the current + `dt_to_t` / `ETP_RULES_DT_TO_T`; two executable cells in + `etp_primitives.ipynb` raised `ImportError` / `TypeError` if re-run. + Incidentally, `etp_primitives.ipynb` and `docs/quickstart/concepts.ipynb` + also still called `element_wise(weight, fn=...)`, predating that + parameter's rename to `weight_fn`. All affected notebooks were re-executed + end-to-end to confirm they now run cleanly. ## Version 0.2.4 diff --git a/docs/advanced/limitations.ipynb b/docs/advanced/limitations.ipynb index c5faf4e..8a0991a 100644 --- a/docs/advanced/limitations.ipynb +++ b/docs/advanced/limitations.ipynb @@ -59,10 +59,10 @@ "id": "a1b2c3d4e5f60005", "metadata": { "execution": { - "iopub.execute_input": "2026-06-26T13:40:58.821878Z", - "iopub.status.busy": "2026-06-26T13:40:58.821727Z", - "iopub.status.idle": "2026-06-26T13:41:05.412843Z", - "shell.execute_reply": "2026-06-26T13:41:05.411595Z" + "iopub.execute_input": "2026-07-09T05:30:50.801846Z", + "iopub.status.busy": "2026-07-09T05:30:50.801678Z", + "iopub.status.idle": "2026-07-09T05:30:52.315086Z", + "shell.execute_reply": "2026-07-09T05:30:52.314133Z" } }, "outputs": [], @@ -116,18 +116,52 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "a1b2c3d4e5f60008", "metadata": { "execution": { - "iopub.execute_input": "2026-06-26T13:41:05.415143Z", - "iopub.status.busy": "2026-06-26T13:41:05.414867Z", - "iopub.status.idle": "2026-06-26T13:41:10.854480Z", - "shell.execute_reply": "2026-06-26T13:41:10.853427Z" + "iopub.execute_input": "2026-07-09T05:30:52.316903Z", + "iopub.status.busy": "2026-07-09T05:30:52.316658Z", + "iopub.status.idle": "2026-07-09T05:30:52.476531Z", + "shell.execute_reply": "2026-07-09T05:30:52.475510Z" } }, - "outputs": [], - "source": "# CORRECT: use jax.lax.select (no sub-Jaxpr) instead of jnp.where (which now jits).\nclass GoodModel(brainstate.nn.Module):\n def __init__(self):\n super().__init__()\n self.w = brainstate.ParamState(jnp.ones((10, 10)))\n self.h = brainstate.HiddenState(jnp.zeros(10))\n\n def update(self, x):\n new_h = jax.nn.tanh(braintrace.matmul(self.h.value, self.w.value) + x)\n # jax.lax.select compiles to a single select_n primitive -- the\n # compiler can trace right through it.\n self.h.value = jax.lax.select(jnp.sum(x) > 0, new_h, self.h.value)\n return self.h.value\n\n\nmodel = GoodModel()\nalgo = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(10), batch_size=1)\nprint(\"Compilation successful.\")" + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Compilation successful.\n" + ] + } + ], + "source": [ + "# CORRECT: use jax.lax.select (no sub-Jaxpr) instead of jnp.where (which now jits).\n", + "class GoodModel(brainstate.nn.Module):\n", + " def __init__(self):\n", + " super().__init__()\n", + " self.w = brainstate.ParamState(jnp.ones((10, 10)))\n", + " self.h = brainstate.HiddenState(jnp.zeros(10))\n", + "\n", + " def update(self, x):\n", + " new_h = jax.nn.tanh(braintrace.matmul(self.h.value, self.w.value) + x)\n", + " # jax.lax.select compiles to a single select_n primitive -- the\n", + " # compiler can trace right through it.\n", + " self.h.value = jax.lax.select(jnp.sum(x) > 0, new_h, self.h.value)\n", + " return self.h.value\n", + "\n", + "\n", + "model = GoodModel()\n", + "algo = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(10), batch_size=1)\n", + "print(\"Compilation successful.\")" + ] }, { "cell_type": "markdown", @@ -164,10 +198,10 @@ "id": "a1b2c3d4e5f60011", "metadata": { "execution": { - "iopub.execute_input": "2026-06-26T13:41:10.856709Z", - "iopub.status.busy": "2026-06-26T13:41:10.856477Z", - "iopub.status.idle": "2026-06-26T13:41:10.860286Z", - "shell.execute_reply": "2026-06-26T13:41:10.859446Z" + "iopub.execute_input": "2026-07-09T05:30:52.478585Z", + "iopub.status.busy": "2026-07-09T05:30:52.478423Z", + "iopub.status.idle": "2026-07-09T05:30:52.482620Z", + "shell.execute_reply": "2026-07-09T05:30:52.481501Z" } }, "outputs": [], @@ -208,7 +242,7 @@ "source": [ "## The \"Weight -> Weight -> Hidden\" Invariant\n", "\n", - "ETP rules are *local* to a single primitive: each primitive's `xy_to_dw` rule assumes its `x` is externally-supplied data, and its `yw_to_w` rule assumes the path from this primitive's output `y` to a hidden state `h` contains **no other trainable ETP weights**. If a primitive `W1`'s output flows through another non-gradient-enabled ETP primitive `W2` before reaching `h`, the per-primitive rules cannot soundly account for `W1` -- that would either double-count the contribution from `W2` or treat `W2`'s `x` as raw data when it is actually a function of `W1`.\n", + "ETP rules are *local* to a single primitive: each primitive's `xy_to_dw` rule assumes its `x` is externally-supplied data, and its `dt_to_t` rule assumes the path from this primitive's output `y` to a hidden state `h` contains **no other trainable ETP weights**. If a primitive `W1`'s output flows through another non-gradient-enabled ETP primitive `W2` before reaching `h`, the per-primitive rules cannot soundly account for `W1` -- that would either double-count the contribution from `W2` or treat `W2`'s `x` as raw data when it is actually a function of `W1`.\n", "\n", "The compiler enforces this by **excluding** `W1` from the relation list whenever its only path to `h` passes through another non-gradient-enabled ETP primitive. The diagnostic kind is `RELATION_EXCLUDED_WEIGHT_TO_WEIGHT` and a `UserWarning` is emitted at compile time. The excluded weight is still trainable -- but only via BPTT, not via online learning.\n", "\n", @@ -233,10 +267,10 @@ "id": "15710888", "metadata": { "execution": { - "iopub.execute_input": "2026-06-26T13:41:10.862181Z", - "iopub.status.busy": "2026-06-26T13:41:10.862018Z", - "iopub.status.idle": "2026-06-26T13:41:11.434153Z", - "shell.execute_reply": "2026-06-26T13:41:11.433146Z" + "iopub.execute_input": "2026-07-09T05:30:52.484962Z", + "iopub.status.busy": "2026-07-09T05:30:52.484764Z", + "iopub.status.idle": "2026-07-09T05:30:53.056501Z", + "shell.execute_reply": "2026-07-09T05:30:53.055618Z" } }, "outputs": [ @@ -257,7 +291,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "/mnt/d/codes/projects/braintrace/braintrace/_compiler/hid_param_op.py:852: UserWarning: ETP primitive etp_mm (weight=('Wr', 'weight')) reaches a hidden state only through another trainable ETP primitive (etp_mm). Per the non-parametric-tail invariant this weight is excluded from ETP; learn it by BPTT or rewire the architecture so its output flows directly into a hidden state.\n", + "/home/chaoming/miniconda3/lib/python3.13/site-packages/braintrace/_compiler/hid_param_op.py:853: UserWarning: ETP primitive etp_mm (weight=('Wr', 'weight')) reaches a hidden state only through another trainable ETP primitive (etp_mm). Per the non-parametric-tail invariant this weight is excluded from ETP; learn it by BPTT or rewire the architecture so its output flows directly into a hidden state.\n", " _emit_no_relation_diag(\n" ] } @@ -330,7 +364,28 @@ "cell_type": "markdown", "id": "a1b2c3d4e5f60014", "metadata": {}, - "source": "## Compilation Time\n\nThe braintrace compiler performs several steps when compiling a model for online learning:\n\n1. **Jaxpr tracing**: JAX traces the model's `update()` method to produce a Jaxpr\n2. **Relation discovery**: The compiler walks the Jaxpr to find ETP primitives, trace weight origins, and connect them to hidden states\n3. **Graph construction**: The eligibility trace computation graph is built from the discovered relations\n\n`braintrace.compile(model, algo, x0, batch_size=B)` runs all these steps and returns a ready learner. This compilation can be slow for complex models, especially on the first call. However:\n\n- **Subsequent calls with the same input shapes reuse the compiled graph.** The compilation result is cached, so you only pay the cost once.\n- **`braintrace.compile` should be called once before the training loop**, not inside it. Calling it repeatedly with the same shapes is harmless (it detects the cache hit), but calling it inside a loop adds unnecessary overhead.\n\n```python\n# Good: compile once, then run many steps\nlearner = braintrace.compile(model, braintrace.D_RTRL, example_input, batch_size=B)\n\nfor step in range(num_steps):\n output = learner(input_data[step]) # uses cached compilation\n```" + "source": [ + "## Compilation Time\n", + "\n", + "The braintrace compiler performs several steps when compiling a model for online learning:\n", + "\n", + "1. **Jaxpr tracing**: JAX traces the model's `update()` method to produce a Jaxpr\n", + "2. **Relation discovery**: The compiler walks the Jaxpr to find ETP primitives, trace weight origins, and connect them to hidden states\n", + "3. **Graph construction**: The eligibility trace computation graph is built from the discovered relations\n", + "\n", + "`braintrace.compile(model, algo, x0, batch_size=B)` runs all these steps and returns a ready learner. This compilation can be slow for complex models, especially on the first call. However:\n", + "\n", + "- **Subsequent calls with the same input shapes reuse the compiled graph.** The compilation result is cached, so you only pay the cost once.\n", + "- **`braintrace.compile` should be called once before the training loop**, not inside it. Calling it repeatedly with the same shapes is harmless (it detects the cache hit), but calling it inside a loop adds unnecessary overhead.\n", + "\n", + "```python\n", + "# Good: compile once, then run many steps\n", + "learner = braintrace.compile(model, braintrace.D_RTRL, example_input, batch_size=B)\n", + "\n", + "for step in range(num_steps):\n", + " output = learner(input_data[step]) # uses cached compilation\n", + "```" + ] }, { "cell_type": "markdown", @@ -378,7 +433,25 @@ "cell_type": "markdown", "id": "a1b2c3d4e5f60016", "metadata": {}, - "source": "## Summary\n\nThe key limitations and their workarounds are:\n\n1. **Avoid `cond`, `scan`, `while_loop`, and nested `vmap` inside the model's `update()` method.** These create sub-Jaxprs that the compiler cannot traverse. Use them freely outside the model (e.g., for time-step unrolling).\n\n2. **Use `jnp.where` and masks as alternatives to conditional logic.** Element-wise selection operations are fully supported and produce equivalent results for most use cases.\n\n3. **Ensure shape compatibility between ETP primitive outputs and hidden states.** The compiler filters connections by shape -- if shapes don't match or broadcast, the connection won't be established.\n\n4. **Per-primitive ETP rules are local.** A weight whose only path to a hidden state passes through another trainable ETP primitive is excluded with a `RELATION_EXCLUDED_WEIGHT_TO_WEIGHT` warning -- it must be learned via BPTT or the architecture must be rewired. `etp_elemwise_p` (the only `gradient_enabled=True` built-in) is the sole exception.\n\n5. **Choose the right algorithm based on memory/accuracy tradeoffs.** Use `D_RTRL` for exact gradients with moderate model sizes, and `ES_D_RTRL` for memory-efficient approximate gradients with larger models.\n\n6. **Call `braintrace.compile` once before training**, not inside the training loop. The compiled graph is cached and reused for inputs of the same shape.\n\n7. **The compiler works with all standard JAX mathematical operations.** As long as you avoid the unsupported control flow primitives listed above, your model will compile successfully." + "source": [ + "## Summary\n", + "\n", + "The key limitations and their workarounds are:\n", + "\n", + "1. **Avoid `cond`, `scan`, `while_loop`, and nested `vmap` inside the model's `update()` method.** These create sub-Jaxprs that the compiler cannot traverse. Use them freely outside the model (e.g., for time-step unrolling).\n", + "\n", + "2. **Use `jnp.where` and masks as alternatives to conditional logic.** Element-wise selection operations are fully supported and produce equivalent results for most use cases.\n", + "\n", + "3. **Ensure shape compatibility between ETP primitive outputs and hidden states.** The compiler filters connections by shape -- if shapes don't match or broadcast, the connection won't be established.\n", + "\n", + "4. **Per-primitive ETP rules are local.** A weight whose only path to a hidden state passes through another trainable ETP primitive is excluded with a `RELATION_EXCLUDED_WEIGHT_TO_WEIGHT` warning -- it must be learned via BPTT or the architecture must be rewired. `etp_elemwise_p` (the only `gradient_enabled=True` built-in) is the sole exception.\n", + "\n", + "5. **Choose the right algorithm based on memory/accuracy tradeoffs.** Use `D_RTRL` for exact gradients with moderate model sizes, and `ES_D_RTRL` for memory-efficient approximate gradients with larger models.\n", + "\n", + "6. **Call `braintrace.compile` once before training**, not inside the training loop. The compiled graph is cached and reused for inputs of the same shape.\n", + "\n", + "7. **The compiler works with all standard JAX mathematical operations.** As long as you avoid the unsupported control flow primitives listed above, your model will compile successfully." + ] } ], "metadata": { @@ -402,4 +475,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/docs/quickstart/concepts.ipynb b/docs/quickstart/concepts.ipynb index 2e78d58..43e085a 100644 --- a/docs/quickstart/concepts.ipynb +++ b/docs/quickstart/concepts.ipynb @@ -101,21 +101,13 @@ "id": "f6a7b8c9d0e1f2a3", "metadata": { "execution": { - "iopub.execute_input": "2026-04-17T09:24:43.456851Z", - "iopub.status.busy": "2026-04-17T09:24:43.456699Z", - "iopub.status.idle": "2026-04-17T09:24:45.629050Z", - "shell.execute_reply": "2026-04-17T09:24:45.628092Z" + "iopub.execute_input": "2026-07-09T05:34:28.616843Z", + "iopub.status.busy": "2026-07-09T05:34:28.616688Z", + "iopub.status.idle": "2026-07-09T05:34:30.919598Z", + "shell.execute_reply": "2026-07-09T05:34:30.918569Z" } }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" - ] - } - ], + "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", @@ -129,10 +121,10 @@ "id": "a7b8c9d0e1f2a3b4", "metadata": { "execution": { - "iopub.execute_input": "2026-04-17T09:24:45.631615Z", - "iopub.status.busy": "2026-04-17T09:24:45.631220Z", - "iopub.status.idle": "2026-04-17T09:24:45.635810Z", - "shell.execute_reply": "2026-04-17T09:24:45.635234Z" + "iopub.execute_input": "2026-07-09T05:34:30.922053Z", + "iopub.status.busy": "2026-07-09T05:34:30.921707Z", + "iopub.status.idle": "2026-07-09T05:34:30.927100Z", + "shell.execute_reply": "2026-07-09T05:34:30.925810Z" } }, "outputs": [], @@ -195,7 +187,7 @@ "For the matching low-level API, the user-facing primitive functions are:\n", "\n", "- `braintrace.matmul(x, w, bias=None)` -- dense matrix multiplication\n", - "- `braintrace.element_wise(weight, fn=...)` -- element-wise weight ops (gating, learnable thresholds)\n", + "- `braintrace.element_wise(weight, weight_fn=...)` -- element-wise weight ops (gating, learnable thresholds)\n", "- `braintrace.conv(x, kernel, bias, ...)` -- convolution\n", "- `braintrace.sparse_matmul(x, weight_data, *, sparse_mat, bias=None)` -- sparse matmul\n", "- `braintrace.lora_matmul(x, B, A, *, alpha=1.0, bias=None)` -- LoRA decomposition\n", @@ -209,10 +201,10 @@ "id": "d0e1f2a3b4c5d6e7", "metadata": { "execution": { - "iopub.execute_input": "2026-04-17T09:24:45.637508Z", - "iopub.status.busy": "2026-04-17T09:24:45.637379Z", - "iopub.status.idle": "2026-04-17T09:24:45.640904Z", - "shell.execute_reply": "2026-04-17T09:24:45.640285Z" + "iopub.execute_input": "2026-07-09T05:34:30.929372Z", + "iopub.status.busy": "2026-07-09T05:34:30.929149Z", + "iopub.status.idle": "2026-07-09T05:34:30.933145Z", + "shell.execute_reply": "2026-07-09T05:34:30.932157Z" } }, "outputs": [], @@ -239,70 +231,174 @@ "cell_type": "markdown", "id": "f2a3b4c5d6e7f8a9", "metadata": {}, - "source": "## 5. Online Learning in 3 Steps\n\nUsing BrainTrace for online learning follows a simple three-step workflow:\n\n1. **Define the model** using `braintrace.nn` modules or manual ETP primitives.\n2. **Compile with `braintrace.compile`** — one call initialises states, builds the eligibility-trace graph, and returns a ready learner.\n3. **Train with standard JAX gradient computation** — eligibility traces are updated inside the wrapped model call.\n\nHere is the complete workflow:" + "source": [ + "## 5. Online Learning in 3 Steps\n", + "\n", + "Using BrainTrace for online learning follows a simple three-step workflow:\n", + "\n", + "1. **Define the model** using `braintrace.nn` modules or manual ETP primitives.\n", + "2. **Compile with `braintrace.compile`** — one call initialises states, builds the eligibility-trace graph, and returns a ready learner.\n", + "3. **Train with standard JAX gradient computation** — eligibility traces are updated inside the wrapped model call.\n", + "\n", + "Here is the complete workflow:" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "a3b4c5d6e7f8a9b0", "metadata": { "execution": { - "iopub.execute_input": "2026-04-17T09:24:45.642700Z", - "iopub.status.busy": "2026-04-17T09:24:45.642477Z", - "iopub.status.idle": "2026-04-17T09:24:47.457511Z", - "shell.execute_reply": "2026-04-17T09:24:47.456529Z" + "iopub.execute_input": "2026-07-09T05:34:30.935096Z", + "iopub.status.busy": "2026-07-09T05:34:30.934843Z", + "iopub.status.idle": "2026-07-09T05:34:32.424103Z", + "shell.execute_reply": "2026-07-09T05:34:32.423104Z" } }, - "outputs": [], - "source": "# Step 1: Define model\nmodel = GRUNet(10, 64, 10)\n\n# Step 2: Compile — initialises states, builds the trace graph, returns a ready learner.\n# The example input carries the batch axis: shape (batch_size, n_in) = (1, 10).\ntrainer = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros((1, 10)), batch_size=1)" + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/mnt/d/codes/projects/braintrace/.claude/worktrees/fix-stale-yw-to-w-docs/braintrace/_compiler/hid_param_op.py:969: UserWarning: ETP primitive etp_mm (weight=('rnn', 'Wr', 'weight')) reaches a hidden state only through another trainable ETP primitive (etp_mm). Per the non-parametric-tail invariant this weight is excluded from ETP; learn it by BPTT or rewire the architecture so its output flows directly into a hidden state.\n", + " _emit_no_relation_diag(\n", + "/mnt/d/codes/projects/braintrace/.claude/worktrees/fix-stale-yw-to-w-docs/braintrace/_compiler/hid_param_op.py:969: UserWarning: ETP primitive etp_mm (weight=('readout', 'weight')) has no connected hidden states. It will be treated as a non-temporal parameter.\n", + " _emit_no_relation_diag(\n" + ] + } + ], + "source": [ + "# Step 1: Define model\n", + "model = GRUNet(10, 64, 10)\n", + "\n", + "# Step 2: Compile — initialises states, builds the trace graph, returns a ready learner.\n", + "# The example input carries the batch axis: shape (batch_size, n_in) = (1, 10).\n", + "trainer = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros((1, 10)), batch_size=1)" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "b4c5d6e7f8a9b0c1", "metadata": { "execution": { - "iopub.execute_input": "2026-04-17T09:24:47.459310Z", - "iopub.status.busy": "2026-04-17T09:24:47.459059Z", - "iopub.status.idle": "2026-04-17T09:24:47.517887Z", - "shell.execute_reply": "2026-04-17T09:24:47.516359Z" + "iopub.execute_input": "2026-07-09T05:34:32.425868Z", + "iopub.status.busy": "2026-07-09T05:34:32.425706Z", + "iopub.status.idle": "2026-07-09T05:34:33.886424Z", + "shell.execute_reply": "2026-07-09T05:34:33.885555Z" } }, "outputs": [], - "source": "# Step 3: Use standard JAX gradient computation\n# The eligibility traces are updated inside trainer(x)\nweights = model.states(brainstate.ParamState)\n\ndef loss_fn(x):\n out = trainer(x)\n return jnp.mean(out ** 2)\n\ngrad_fn = brainstate.transform.grad(loss_fn, weights)\ngrads = grad_fn(jnp.ones((1, 10))) # input carries the batch axis (1, 10)" + "source": [ + "# Step 3: Use standard JAX gradient computation\n", + "# The eligibility traces are updated inside trainer(x)\n", + "weights = model.states(brainstate.ParamState)\n", + "\n", + "def loss_fn(x):\n", + " out = trainer(x)\n", + " return jnp.mean(out ** 2)\n", + "\n", + "grad_fn = brainstate.transform.grad(loss_fn, weights)\n", + "grads = grad_fn(jnp.ones((1, 10))) # input carries the batch axis (1, 10)" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, "id": "c5d6e7f8a9b0c1d2", "metadata": { "execution": { - "iopub.execute_input": "2026-04-17T09:24:47.520581Z", - "iopub.status.busy": "2026-04-17T09:24:47.520313Z", - "iopub.status.idle": "2026-04-17T09:24:47.341332Z", - "shell.execute_reply": "2026-04-17T09:24:47.340466Z" + "iopub.execute_input": "2026-07-09T05:34:33.888453Z", + "iopub.status.busy": "2026-07-09T05:34:33.888294Z", + "iopub.status.idle": "2026-07-09T05:34:34.163400Z", + "shell.execute_reply": "2026-07-09T05:34:34.162356Z" } }, "outputs": [], - "source": "# Step 3: Use standard JAX gradient computation\n# The eligibility traces are updated inside trainer(x)\nweights = model.states(brainstate.ParamState)\n\ndef loss_fn(x):\n out = trainer(x)\n return jnp.mean(out ** 2)\n\ngrad_fn = brainstate.transform.grad(loss_fn, weights)\ngrads = grad_fn(jnp.ones((1, 10))) # input carries the batch axis (1, 10)" + "source": [ + "# Step 3: Use standard JAX gradient computation\n", + "# The eligibility traces are updated inside trainer(x)\n", + "weights = model.states(brainstate.ParamState)\n", + "\n", + "def loss_fn(x):\n", + " out = trainer(x)\n", + " return jnp.mean(out ** 2)\n", + "\n", + "grad_fn = brainstate.transform.grad(loss_fn, weights)\n", + "grads = grad_fn(jnp.ones((1, 10))) # input carries the batch axis (1, 10)" + ] }, { "cell_type": "markdown", "id": "d6e7f8a9b0c1d2e3", "metadata": {}, - "source": "**What happens under the hood:**\n\n- `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=1)` initialises all model states, traces the model through JAX, identifies all ETP primitives, and builds the eligibility trace computation graph — all in one call.\n- Each call to `trainer(x)` runs the model forward pass *and* updates all eligibility traces.\n- When you compute `grad(loss_fn, weights)`, the algorithm uses `custom_vjp` to provide gradients that incorporate the eligibility trace information — giving you temporally-aware gradients with $O(1)$ memory per step." + "source": [ + "**What happens under the hood:**\n", + "\n", + "- `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=1)` initialises all model states, traces the model through JAX, identifies all ETP primitives, and builds the eligibility trace computation graph — all in one call.\n", + "- Each call to `trainer(x)` runs the model forward pass *and* updates all eligibility traces.\n", + "- When you compute `grad(loss_fn, weights)`, the algorithm uses `custom_vjp` to provide gradients that incorporate the eligibility trace information — giving you temporally-aware gradients with $O(1)$ memory per step." + ] }, { "cell_type": "markdown", "id": "e7f8a9b0c1d2e3f4", "metadata": {}, - "source": "## 6. Available Algorithms\n\nBrainTrace provides two main online learning algorithms:\n\n| Algorithm | Class | Memory Complexity | Compute Complexity | Best For |\n|---|---|---|---|---|\n| **D-RTRL** | `braintrace.D_RTRL` | $O(B \\cdot \\theta)$ | $O(B \\cdot I \\cdot O)$ per layer | RNNs (GRU, LSTM) with moderate hidden sizes |\n| **ES-D-RTRL** | `braintrace.ES_D_RTRL` | $O(B \\cdot N)$ | $O(B \\cdot N)$ per layer | Large-scale SNNs where $N$ is the number of neurons |\n\nWhere $B$ is the batch size, $\\theta$ is the number of parameters, $I$ and $O$ are input/output dimensions, and $N$ is the number of neurons.\n\n### When to use which?\n\n- **`D_RTRL`** (also `ParamDimVjpAlgorithm`): Use for **rate-based RNNs** (GRU, LSTM, etc.) where you need accurate temporal gradient propagation. Its $O(\\theta)$ memory cost scales with parameter count, which is acceptable for typical RNN hidden sizes.\n\n- **`pp_prop`** (also `ES_D_RTRL` / `IODimVjpAlgorithm`): Use for **spiking neural networks** (SNNs) or very large recurrent networks. It achieves $O(N)$ complexity by exploiting the element-wise nature of neuronal dynamics, making it much more efficient for networks with many neurons.\n\nBoth algorithms are accessed through `braintrace.compile` — just pass the class name:\n\n```python\n# D-RTRL for RNNs\nlearner = braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B)\n\n# ES-D-RTRL for SNNs (pass algorithm hyperparameters as kwargs)\nlearner = braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, decay_or_rank=0.99)\n```" + "source": [ + "## 6. Available Algorithms\n", + "\n", + "BrainTrace provides two main online learning algorithms:\n", + "\n", + "| Algorithm | Class | Memory Complexity | Compute Complexity | Best For |\n", + "|---|---|---|---|---|\n", + "| **D-RTRL** | `braintrace.D_RTRL` | $O(B \\cdot \\theta)$ | $O(B \\cdot I \\cdot O)$ per layer | RNNs (GRU, LSTM) with moderate hidden sizes |\n", + "| **ES-D-RTRL** | `braintrace.ES_D_RTRL` | $O(B \\cdot N)$ | $O(B \\cdot N)$ per layer | Large-scale SNNs where $N$ is the number of neurons |\n", + "\n", + "Where $B$ is the batch size, $\\theta$ is the number of parameters, $I$ and $O$ are input/output dimensions, and $N$ is the number of neurons.\n", + "\n", + "### When to use which?\n", + "\n", + "- **`D_RTRL`** (also `ParamDimVjpAlgorithm`): Use for **rate-based RNNs** (GRU, LSTM, etc.) where you need accurate temporal gradient propagation. Its $O(\\theta)$ memory cost scales with parameter count, which is acceptable for typical RNN hidden sizes.\n", + "\n", + "- **`pp_prop`** (also `ES_D_RTRL` / `IODimVjpAlgorithm`): Use for **spiking neural networks** (SNNs) or very large recurrent networks. It achieves $O(N)$ complexity by exploiting the element-wise nature of neuronal dynamics, making it much more efficient for networks with many neurons.\n", + "\n", + "Both algorithms are accessed through `braintrace.compile` — just pass the class name:\n", + "\n", + "```python\n", + "# D-RTRL for RNNs\n", + "learner = braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B)\n", + "\n", + "# ES-D-RTRL for SNNs (pass algorithm hyperparameters as kwargs)\n", + "learner = braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, decay_or_rank=0.99)\n", + "```" + ] }, { "cell_type": "markdown", "id": "f8a9b0c1d2e3f4a5", "metadata": {}, - "source": "## 7. Summary\n\nHere is a quick recap of the core concepts:\n\n| Concept | Description |\n|---|---|\n| **Online learning** | Update weights at each time step using eligibility traces, achieving $O(1)$ memory per step |\n| **ETP primitives** | `braintrace.matmul`, `braintrace.element_wise`, `braintrace.conv` — JAX custom primitives that mark operations for online learning |\n| **Primitive-based selection** | Use an ETP primitive to include a weight; use regular JAX ops to exclude it |\n| **`braintrace.nn`** | Pre-built layers (Linear, GRUCell, LSTMCell, Conv) that use ETP primitives internally |\n| **`braintrace.compile`** | `braintrace.compile(model, algo, x0, batch_size=B)` — the one-call entry point that initialises states, builds the trace graph, and returns a ready learner |\n| **D_RTRL** | $O(\\theta)$ algorithm for RNNs |\n| **ES_D_RTRL** | $O(N)$ algorithm for SNNs |" + "source": [ + "## 7. Summary\n", + "\n", + "Here is a quick recap of the core concepts:\n", + "\n", + "| Concept | Description |\n", + "|---|---|\n", + "| **Online learning** | Update weights at each time step using eligibility traces, achieving $O(1)$ memory per step |\n", + "| **ETP primitives** | `braintrace.matmul`, `braintrace.element_wise`, `braintrace.conv` — JAX custom primitives that mark operations for online learning |\n", + "| **Primitive-based selection** | Use an ETP primitive to include a weight; use regular JAX ops to exclude it |\n", + "| **`braintrace.nn`** | Pre-built layers (Linear, GRUCell, LSTMCell, Conv) that use ETP primitives internally |\n", + "| **`braintrace.compile`** | `braintrace.compile(model, algo, x0, batch_size=B)` — the one-call entry point that initialises states, builds the trace graph, and returns a ready learner |\n", + "| **D_RTRL** | $O(\\theta)$ algorithm for RNNs |\n", + "| **ES_D_RTRL** | $O(N)$ algorithm for SNNs |" + ] }, { "cell_type": "markdown", @@ -342,4 +438,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/docs/tutorials/customizing_primitive_transforms.ipynb b/docs/tutorials/customizing_primitive_transforms.ipynb index c1ae5fd..7d10a78 100644 --- a/docs/tutorials/customizing_primitive_transforms.ipynb +++ b/docs/tutorials/customizing_primitive_transforms.ipynb @@ -26,7 +26,7 @@ "\n", "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$.\n", "\n", - "**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.\n", + "**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 (`dt_to_t`, `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.\n", "\n", "$$\n", "\\frac{\\partial h}{\\partial W_{\\text{raw}}}\n", @@ -43,10 +43,10 @@ "id": "6387391b", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:09:38.724504Z", - "iopub.status.busy": "2026-06-30T04:09:38.724361Z", - "iopub.status.idle": "2026-06-30T04:09:41.868101Z", - "shell.execute_reply": "2026-06-30T04:09:41.867212Z" + "iopub.execute_input": "2026-07-09T05:32:06.898364Z", + "iopub.status.busy": "2026-07-09T05:32:06.898115Z", + "iopub.status.idle": "2026-07-09T05:32:09.047145Z", + "shell.execute_reply": "2026-07-09T05:32:09.045766Z" } }, "outputs": [], @@ -87,13 +87,20 @@ "id": "fc2ffd41", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:09:41.870171Z", - "iopub.status.busy": "2026-06-30T04:09:41.869945Z", - "iopub.status.idle": "2026-06-30T04:10:04.109078Z", - "shell.execute_reply": "2026-06-30T04:10:04.107531Z" + "iopub.execute_input": "2026-07-09T05:32:09.049255Z", + "iopub.status.busy": "2026-07-09T05:32:09.048995Z", + "iopub.status.idle": "2026-07-09T05:32:09.704196Z", + "shell.execute_reply": "2026-07-09T05:32:09.703229Z" } }, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -119,10 +126,10 @@ "id": "04f46b2f", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:04.113648Z", - "iopub.status.busy": "2026-06-30T04:10:04.113360Z", - "iopub.status.idle": "2026-06-30T04:10:17.729431Z", - "shell.execute_reply": "2026-06-30T04:10:17.727833Z" + "iopub.execute_input": "2026-07-09T05:32:09.706855Z", + "iopub.status.busy": "2026-07-09T05:32:09.706685Z", + "iopub.status.idle": "2026-07-09T05:32:11.099348Z", + "shell.execute_reply": "2026-07-09T05:32:11.098258Z" } }, "outputs": [ @@ -176,10 +183,10 @@ "id": "9c2d84de", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:17.732106Z", - "iopub.status.busy": "2026-06-30T04:10:17.731847Z", - "iopub.status.idle": "2026-06-30T04:10:17.868021Z", - "shell.execute_reply": "2026-06-30T04:10:17.865725Z" + "iopub.execute_input": "2026-07-09T05:32:11.101376Z", + "iopub.status.busy": "2026-07-09T05:32:11.101203Z", + "iopub.status.idle": "2026-07-09T05:32:11.182300Z", + "shell.execute_reply": "2026-07-09T05:32:11.181050Z" } }, "outputs": [ @@ -222,10 +229,10 @@ "id": "1791e74c", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:17.872301Z", - "iopub.status.busy": "2026-06-30T04:10:17.871984Z", - "iopub.status.idle": "2026-06-30T04:10:31.712950Z", - "shell.execute_reply": "2026-06-30T04:10:31.711941Z" + "iopub.execute_input": "2026-07-09T05:32:11.184225Z", + "iopub.status.busy": "2026-07-09T05:32:11.184066Z", + "iopub.status.idle": "2026-07-09T05:32:12.632809Z", + "shell.execute_reply": "2026-07-09T05:32:12.631801Z" } }, "outputs": [ @@ -301,10 +308,10 @@ "id": "30ae0189", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:31.715494Z", - "iopub.status.busy": "2026-06-30T04:10:31.715271Z", - "iopub.status.idle": "2026-06-30T04:10:31.719393Z", - "shell.execute_reply": "2026-06-30T04:10:31.718301Z" + "iopub.execute_input": "2026-07-09T05:32:12.634573Z", + "iopub.status.busy": "2026-07-09T05:32:12.634401Z", + "iopub.status.idle": "2026-07-09T05:32:12.638197Z", + "shell.execute_reply": "2026-07-09T05:32:12.637330Z" } }, "outputs": [], @@ -348,10 +355,10 @@ "id": "508ef4c6", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:31.721023Z", - "iopub.status.busy": "2026-06-30T04:10:31.720890Z", - "iopub.status.idle": "2026-06-30T04:10:31.726327Z", - "shell.execute_reply": "2026-06-30T04:10:31.725521Z" + "iopub.execute_input": "2026-07-09T05:32:12.639984Z", + "iopub.status.busy": "2026-07-09T05:32:12.639782Z", + "iopub.status.idle": "2026-07-09T05:32:12.646762Z", + "shell.execute_reply": "2026-07-09T05:32:12.645848Z" } }, "outputs": [ @@ -364,10 +371,10 @@ } ], "source": [ - "# Step 2: the four ETP rules. Note: yw_to_w / init_* are TRANSFORM-FREE.\n", + "# Step 2: the four ETP rules. Note: dt_to_t / init_* are TRANSFORM-FREE.\n", "# Only xy_to_dw differentiates the impl (which applies f), so f' enters here alone.\n", "\n", - "def _scaled_yw_to_w(hidden_dim, trace, *, scale=1.0, has_bias=False,\n", + "def _scaled_dt_to_t(hidden_dim, trace, *, scale=1.0, has_bias=False,\n", " weight_fn=None, bias_fn=None):\n", " # y -> W chain link for y = scale * x @ W: scale on the matching out column.\n", " # No f' here by design.\n", @@ -410,7 +417,7 @@ "\n", "\n", "scaled_mm_p.register_etp_rules(\n", - " yw_to_w=_scaled_yw_to_w,\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", @@ -424,10 +431,10 @@ "id": "2dc77b4c", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:31.727866Z", - "iopub.status.busy": "2026-06-30T04:10:31.727727Z", - "iopub.status.idle": "2026-06-30T04:10:38.044842Z", - "shell.execute_reply": "2026-06-30T04:10:38.043871Z" + "iopub.execute_input": "2026-07-09T05:32:12.648377Z", + "iopub.status.busy": "2026-07-09T05:32:12.648221Z", + "iopub.status.idle": "2026-07-09T05:32:12.941837Z", + "shell.execute_reply": "2026-07-09T05:32:12.940780Z" } }, "outputs": [ @@ -467,10 +474,10 @@ "id": "92a7a949", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:38.047374Z", - "iopub.status.busy": "2026-06-30T04:10:38.047106Z", - "iopub.status.idle": "2026-06-30T04:10:38.753482Z", - "shell.execute_reply": "2026-06-30T04:10:38.752496Z" + "iopub.execute_input": "2026-07-09T05:32:12.943862Z", + "iopub.status.busy": "2026-07-09T05:32:12.943691Z", + "iopub.status.idle": "2026-07-09T05:32:13.337321Z", + "shell.execute_reply": "2026-07-09T05:32:13.336461Z" } }, "outputs": [ @@ -529,7 +536,7 @@ "source": [ "## 5. (Advanced) The fast path and its transform-aware gate\n", "\n", - "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)`.\n", + "For the elementwise-`dt_to_t` 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)`.\n", "\n", "- `instant` — the instantaneous $\\operatorname{diag}(\\mathbf{D}_f^t) \\otimes \\mathbf{x}^t$ term.\n", "- `recurrent` — the $\\mathbf{D}^t \\boldsymbol{\\epsilon}^{t-1}$ contraction.\n", @@ -547,10 +554,10 @@ "id": "1bf1f341", "metadata": { "execution": { - "iopub.execute_input": "2026-06-30T04:10:38.755192Z", - "iopub.status.busy": "2026-06-30T04:10:38.755031Z", - "iopub.status.idle": "2026-06-30T04:10:38.759114Z", - "shell.execute_reply": "2026-06-30T04:10:38.758302Z" + "iopub.execute_input": "2026-07-09T05:32:13.339134Z", + "iopub.status.busy": "2026-07-09T05:32:13.338993Z", + "iopub.status.idle": "2026-07-09T05:32:13.343440Z", + "shell.execute_reply": "2026-07-09T05:32:13.342133Z" } }, "outputs": [ @@ -599,7 +606,7 @@ "## Summary\n", "\n", "- **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.\n", - "- **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.\n", + "- **One rule:** the transform Jacobian $f'$ enters *only* in `xy_to_dw` (via `jax.vjp` through the impl). `dt_to_t` and the trace initializers stay transform-free and exact.\n", "- **Exactness:** `D_RTRL == BPTT` element-wise even with `weight_fn=tanh`, for both built-in and custom primitives — verified via `braintrace._algorithm.oracle`.\n", "- **Custom primitives** need no rule changes for transforms: apply the hook in the impl, and `xy_to_dw`'s VJP threads $f'$ for free.\n", "- **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." diff --git a/docs/tutorials/etp_primitives.ipynb b/docs/tutorials/etp_primitives.ipynb index bbbee06..2cec862 100644 --- a/docs/tutorials/etp_primitives.ipynb +++ b/docs/tutorials/etp_primitives.ipynb @@ -24,7 +24,7 @@ "- **All JAX rules are auto-derived.** Each primitive's `impl` delegates to standard JAX ops (e.g., `x @ w`, `jax.lax.conv_general_dilated`). The `register_primitive()` helper automatically derives `abstract_eval`, MLIR lowering, JVP, transpose, and batching rules from the implementation. No hand-written derivative formulas are needed.\n", "\n", "- **Only ETP-specific rules need hand-writing.** Four global dictionaries capture the online-learning-specific logic:\n", - " - `ETP_RULES_YW_TO_W` — D-RTRL trace propagation (the $\\mathbf{D}^t \\boldsymbol{\\epsilon}^{t-1}$ term)\n", + " - `ETP_RULES_DT_TO_T` — D-RTRL trace propagation (the $\\mathbf{D}^t \\boldsymbol{\\epsilon}^{t-1}$ term)\n", " - `ETP_RULES_XY_TO_DW` — instantaneous hidden-to-weight Jacobian (the $\\operatorname{diag}(\\mathbf{D}_f^t) \\otimes \\mathbf{x}^t$ term)\n", " - `ETP_RULES_INIT_DRTRL` — D-RTRL parameter-dimension trace initialiser\n", " - `ETP_RULES_INIT_PP` — pp-prop / ES-D-RTRL output-dimension df-trace initialiser\n", @@ -40,21 +40,13 @@ "id": "c3d4e5f6", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:53.778568Z", - "iopub.status.busy": "2026-04-18T06:22:53.778250Z", - "iopub.status.idle": "2026-04-18T06:22:56.070543Z", - "shell.execute_reply": "2026-04-18T06:22:56.069694Z" + "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": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" - ] - } - ], + "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", @@ -102,13 +94,20 @@ "id": "f6a7b8c9", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:56.073864Z", - "iopub.status.busy": "2026-04-18T06:22:56.073334Z", - "iopub.status.idle": "2026-04-18T06:22:56.252956Z", - "shell.execute_reply": "2026-04-18T06:22:56.252077Z" + "iopub.execute_input": "2026-07-09T05:33:32.320214Z", + "iopub.status.busy": "2026-07-09T05:33:32.319961Z", + "iopub.status.idle": "2026-07-09T05:33:32.525704Z", + "shell.execute_reply": "2026-07-09T05:33:32.524759Z" } }, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -139,10 +138,10 @@ "id": "a7b8c9d0", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:56.255033Z", - "iopub.status.busy": "2026-04-18T06:22:56.254794Z", - "iopub.status.idle": "2026-04-18T06:22:56.307499Z", - "shell.execute_reply": "2026-04-18T06:22:56.306626Z" + "iopub.execute_input": "2026-07-09T05:33:32.528326Z", + "iopub.status.busy": "2026-07-09T05:33:32.528149Z", + "iopub.status.idle": "2026-07-09T05:33:32.582644Z", + "shell.execute_reply": "2026-07-09T05:33:32.581735Z" } }, "outputs": [ @@ -167,13 +166,13 @@ "id": "b8c9d0e1", "metadata": {}, "source": [ - "### 2. `braintrace.element_wise(weight, fn=lambda w: w)` — Element-wise Operation\n", + "### 2. `braintrace.element_wise(weight, *, weight_fn=None)` — Element-wise Operation\n", "\n", - "Applies `fn` to the weight and passes the result through a marker primitive. The operation is treated as *diagonal* in the hidden-state space.\n", + "Applies `weight_fn` to the weight and passes the result through a marker primitive. The operation is treated as *diagonal* in the hidden-state space.\n", "\n", - "$$y = \\texttt{fn}(w)$$\n", + "$$y = \\texttt{weight\\_fn}(w)$$\n", "\n", - "`fn` defaults to the identity (`lambda w: w`); supply any JAX-differentiable function when you want a non-trivial transformation.\n", + "`weight_fn` defaults to `None` (identity); supply any JAX-differentiable function when you want a non-trivial transformation.\n", "\n", "Common use cases:\n", "- Gating mechanisms in RNNs (learnable gate biases)\n", @@ -189,10 +188,10 @@ "id": "c9d0e1f2", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:56.311150Z", - "iopub.status.busy": "2026-04-18T06:22:56.310960Z", - "iopub.status.idle": "2026-04-18T06:22:56.379757Z", - "shell.execute_reply": "2026-04-18T06:22:56.378983Z" + "iopub.execute_input": "2026-07-09T05:33:32.584737Z", + "iopub.status.busy": "2026-07-09T05:33:32.584433Z", + "iopub.status.idle": "2026-07-09T05:33:32.651206Z", + "shell.execute_reply": "2026-07-09T05:33:32.650196Z" } }, "outputs": [ @@ -207,18 +206,18 @@ } ], "source": [ - "# Identity (default fn): just marks the weight for ETP\n", + "# Identity (default weight_fn): just marks the weight for ETP\n", "w_gate = jnp.array([0.5, -0.3, 0.8, 0.1])\n", "\n", "y_identity = braintrace.element_wise(w_gate)\n", "print(\"Identity:\", y_identity)\n", "\n", "# With a transformation function\n", - "y_sigmoid = braintrace.element_wise(w_gate, fn=jax.nn.sigmoid)\n", + "y_sigmoid = braintrace.element_wise(w_gate, weight_fn=jax.nn.sigmoid)\n", "print(\"Sigmoid:\", y_sigmoid)\n", "\n", "# With absolute value (e.g., enforcing positive time constants)\n", - "y_abs = braintrace.element_wise(w_gate, fn=jnp.abs)\n", + "y_abs = braintrace.element_wise(w_gate, weight_fn=jnp.abs)\n", "print(\"Abs:\", y_abs)" ] }, @@ -244,10 +243,10 @@ "id": "e1f2a3b4", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:56.381550Z", - "iopub.status.busy": "2026-04-18T06:22:56.381398Z", - "iopub.status.idle": "2026-04-18T06:22:56.479299Z", - "shell.execute_reply": "2026-04-18T06:22:56.478311Z" + "iopub.execute_input": "2026-07-09T05:33:32.653207Z", + "iopub.status.busy": "2026-07-09T05:33:32.653040Z", + "iopub.status.idle": "2026-07-09T05:33:32.741742Z", + "shell.execute_reply": "2026-07-09T05:33:32.740889Z" } }, "outputs": [ @@ -280,10 +279,10 @@ "id": "f2a3b4c5", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:56.481773Z", - "iopub.status.busy": "2026-04-18T06:22:56.481499Z", - "iopub.status.idle": "2026-04-18T06:22:56.587670Z", - "shell.execute_reply": "2026-04-18T06:22:56.587045Z" + "iopub.execute_input": "2026-07-09T05:33:32.743904Z", + "iopub.status.busy": "2026-07-09T05:33:32.743708Z", + "iopub.status.idle": "2026-07-09T05:33:32.833682Z", + "shell.execute_reply": "2026-07-09T05:33:32.832764Z" } }, "outputs": [ @@ -325,18 +324,46 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "id": "b4c5d6e7", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:56.589616Z", - "iopub.status.busy": "2026-04-18T06:22:56.589449Z", - "iopub.status.idle": "2026-04-18T06:22:57.899681Z", - "shell.execute_reply": "2026-04-18T06:22:57.898623Z" + "iopub.execute_input": "2026-07-09T05:33:32.835474Z", + "iopub.status.busy": "2026-07-09T05:33:32.835337Z", + "iopub.status.idle": "2026-07-09T05:33:35.770728Z", + "shell.execute_reply": "2026-07-09T05:33:35.769886Z" } }, - "outputs": [], - "source": "import brainunit as u\nimport brainevent\n\n# Create a sparse connectivity matrix\ndense_w = jnp.where(\n jax.random.uniform(jax.random.PRNGKey(0), (50, 50)) < 0.1,\n jax.random.normal(jax.random.PRNGKey(1), (50, 50)),\n 0.0\n)\n# sparse_mat must be a brainevent.DataRepresentation (e.g. brainevent.CSR),\n# which implements the with_data / dt2t_transposed / dt2t protocol.\nsparse_mat = brainevent.CSR.fromdense(dense_w)\n\n# The learnable parameter is just the non-zero data\nweight_data = sparse_mat.data\n\nx_sp = jnp.ones((4, 50)) # batch=4, features=50\ny_sp = braintrace.sparse_matmul(x_sp, weight_data, sparse_mat=sparse_mat)\nprint(\"Sparse matmul output shape:\", y_sp.shape) # (4, 50)" + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Sparse matmul output shape: (4, 50)\n" + ] + } + ], + "source": [ + "import brainunit as u\n", + "import brainevent\n", + "\n", + "# Create a sparse connectivity matrix\n", + "dense_w = jnp.where(\n", + " jax.random.uniform(jax.random.PRNGKey(0), (50, 50)) < 0.1,\n", + " jax.random.normal(jax.random.PRNGKey(1), (50, 50)),\n", + " 0.0\n", + ")\n", + "# sparse_mat must be a brainevent.DataRepresentation (e.g. brainevent.CSR),\n", + "# which implements the with_data / dt2t_transposed / dt2t protocol.\n", + "sparse_mat = brainevent.CSR.fromdense(dense_w)\n", + "\n", + "# The learnable parameter is just the non-zero data\n", + "weight_data = sparse_mat.data\n", + "\n", + "x_sp = jnp.ones((4, 50)) # batch=4, features=50\n", + "y_sp = braintrace.sparse_matmul(x_sp, weight_data, sparse_mat=sparse_mat)\n", + "print(\"Sparse matmul output shape:\", y_sp.shape) # (4, 50)" + ] }, { "cell_type": "markdown", @@ -358,10 +385,10 @@ "id": "d6e7f8a9", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:57.901964Z", - "iopub.status.busy": "2026-04-18T06:22:57.901698Z", - "iopub.status.idle": "2026-04-18T06:22:58.464640Z", - "shell.execute_reply": "2026-04-18T06:22:58.463547Z" + "iopub.execute_input": "2026-07-09T05:33:35.773052Z", + "iopub.status.busy": "2026-07-09T05:33:35.772887Z", + "iopub.status.idle": "2026-07-09T05:33:36.321372Z", + "shell.execute_reply": "2026-07-09T05:33:36.320482Z" } }, "outputs": [ @@ -370,14 +397,14 @@ "output_type": "stream", "text": [ "LoRA output shape: (8, 32)\n", - "LoRA output (first sample): [ 1.2993842e-04 2.9886682e-03 -5.7171844e-04 1.7000248e-03\n", - " 2.8961061e-03 -1.4734608e-03 1.5826351e-03 1.1692893e-04\n", - " -5.4561032e-04 -1.3646147e-03 1.0654312e-04 2.0160272e-03\n", - " 3.1371990e-03 1.8710974e-03 2.9124888e-03 -8.2047645e-04\n", - " 9.7736251e-04 -1.1875220e-03 -2.1796541e-03 -5.8191392e-05\n", - " -1.5415263e-03 -1.3381819e-03 -2.1044153e-03 -8.4472442e-04\n", - " 1.6430757e-05 -2.9564608e-04 4.5550600e-04 -1.6011632e-03\n", - " -1.4516843e-03 -1.3209208e-03 3.2850087e-04 -6.0276774e-04]\n" + "LoRA output (first sample): [ 1.29938402e-04 2.98866816e-03 -5.71718439e-04 1.70002470e-03\n", + " 2.89610587e-03 -1.47346081e-03 1.58263510e-03 1.16928706e-04\n", + " -5.45610383e-04 -1.36461481e-03 1.06543179e-04 2.01602723e-03\n", + " 3.13719874e-03 1.87109748e-03 2.91248900e-03 -8.20476678e-04\n", + " 9.77362506e-04 -1.18752208e-03 -2.17965432e-03 -5.81911918e-05\n", + " -1.54152629e-03 -1.33818167e-03 -2.10441509e-03 -8.44724302e-04\n", + " 1.64309131e-05 -2.95645994e-04 4.55506030e-04 -1.60116307e-03\n", + " -1.45168405e-03 -1.32092077e-03 3.28500755e-04 -6.02767745e-04]\n" ] } ], @@ -416,10 +443,10 @@ "id": "1b2c54cf", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.466411Z", - "iopub.status.busy": "2026-04-18T06:22:58.466259Z", - "iopub.status.idle": "2026-04-18T06:22:58.509627Z", - "shell.execute_reply": "2026-04-18T06:22:58.508968Z" + "iopub.execute_input": "2026-07-09T05:33:36.323365Z", + "iopub.status.busy": "2026-07-09T05:33:36.323086Z", + "iopub.status.idle": "2026-07-09T05:33:36.359191Z", + "shell.execute_reply": "2026-07-09T05:33:36.358273Z" } }, "outputs": [ @@ -464,10 +491,10 @@ "id": "f8a9b0c1", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.512275Z", - "iopub.status.busy": "2026-04-18T06:22:58.511982Z", - "iopub.status.idle": "2026-04-18T06:22:58.539485Z", - "shell.execute_reply": "2026-04-18T06:22:58.538856Z" + "iopub.execute_input": "2026-07-09T05:33:36.361884Z", + "iopub.status.busy": "2026-07-09T05:33:36.361676Z", + "iopub.status.idle": "2026-07-09T05:33:36.385708Z", + "shell.execute_reply": "2026-07-09T05:33:36.384820Z" } }, "outputs": [ @@ -494,10 +521,10 @@ "id": "a9b0c1d2", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.541162Z", - "iopub.status.busy": "2026-04-18T06:22:58.541001Z", - "iopub.status.idle": "2026-04-18T06:22:58.713922Z", - "shell.execute_reply": "2026-04-18T06:22:58.713110Z" + "iopub.execute_input": "2026-07-09T05:33:36.387510Z", + "iopub.status.busy": "2026-07-09T05:33:36.387355Z", + "iopub.status.idle": "2026-07-09T05:33:36.575134Z", + "shell.execute_reply": "2026-07-09T05:33:36.574189Z" } }, "outputs": [ @@ -527,10 +554,10 @@ "id": "b0c1d2e3", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.715971Z", - "iopub.status.busy": "2026-04-18T06:22:58.715797Z", - "iopub.status.idle": "2026-04-18T06:22:58.796595Z", - "shell.execute_reply": "2026-04-18T06:22:58.795935Z" + "iopub.execute_input": "2026-07-09T05:33:36.578476Z", + "iopub.status.busy": "2026-07-09T05:33:36.578199Z", + "iopub.status.idle": "2026-07-09T05:33:36.651220Z", + "shell.execute_reply": "2026-07-09T05:33:36.650466Z" } }, "outputs": [ @@ -540,6 +567,14 @@ "text": [ "vmap output shape: (8, 4, 5)\n" ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/chaoming/miniconda3/lib/python3.13/site-packages/jax/_src/interpreters/batching.py:603: UserWarning: ETP primitive 'etp_mm' was decomposed into standard JAX ops under vmap (batch dims (0, None)). Its trainable parameters will NOT be recognized by the eligibility-trace compiler if this trace is compiled for online learning. Map over the data input only (weights unbatched), or call the batched op directly.\n", + " return batcher(vals, dims, **params)\n" + ] } ], "source": [ @@ -557,10 +592,10 @@ "id": "c1d2e3f4", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.798370Z", - "iopub.status.busy": "2026-04-18T06:22:58.798129Z", - "iopub.status.idle": "2026-04-18T06:22:58.866413Z", - "shell.execute_reply": "2026-04-18T06:22:58.865537Z" + "iopub.execute_input": "2026-07-09T05:33:36.653251Z", + "iopub.status.busy": "2026-07-09T05:33:36.653079Z", + "iopub.status.idle": "2026-07-09T05:33:36.707484Z", + "shell.execute_reply": "2026-07-09T05:33:36.706614Z" } }, "outputs": [ @@ -589,10 +624,10 @@ "id": "d2e3f4a5", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.868117Z", - "iopub.status.busy": "2026-04-18T06:22:58.867904Z", - "iopub.status.idle": "2026-04-18T06:22:58.929017Z", - "shell.execute_reply": "2026-04-18T06:22:58.928061Z" + "iopub.execute_input": "2026-07-09T05:33:36.709443Z", + "iopub.status.busy": "2026-07-09T05:33:36.709247Z", + "iopub.status.idle": "2026-07-09T05:33:36.759603Z", + "shell.execute_reply": "2026-07-09T05:33:36.758480Z" } }, "outputs": [ @@ -622,7 +657,46 @@ "cell_type": "markdown", "id": "e3f4a5b6", "metadata": {}, - "source": "## Argument Conventions\n\nEvery 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\nInstead of hard-coding a single weight index, each primitive registers a function\n\n```python\ntrainable_invars_fn: Callable[[dict], Dict[str, int]]\n```\n\nwhich 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\nBuilt-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\nNotes:\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." + "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", @@ -635,10 +709,10 @@ "\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_YW_TO_W` — D-RTRL trace propagation\n", + "### `ETP_RULES_DT_TO_T` — D-RTRL trace propagation\n", "\n", "```python\n", - "yw_to_w(hidden_dim: Array, trace: Dict[str, Array], **static_params) -> Dict[str, Array]\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", @@ -676,10 +750,10 @@ "id": "a5b6c7d8", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.931264Z", - "iopub.status.busy": "2026-04-18T06:22:58.931098Z", - "iopub.status.idle": "2026-04-18T06:22:58.937585Z", - "shell.execute_reply": "2026-04-18T06:22:58.936469Z" + "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": [ @@ -689,7 +763,12 @@ "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", @@ -697,9 +776,14 @@ " etp_sp_mm [batched]\n", " etp_sp_mv\n", "\n", - "Trace propagation rules (ETP_RULES_YW_TO_W):\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", @@ -709,7 +793,12 @@ "\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", @@ -719,7 +808,12 @@ "\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", @@ -729,7 +823,12 @@ "\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", @@ -741,7 +840,7 @@ ], "source": [ "from braintrace._op import (\n", - " ETP_RULES_YW_TO_W,\n", + " ETP_RULES_DT_TO_T,\n", " ETP_RULES_XY_TO_DW,\n", " ETP_RULES_INIT_DRTRL,\n", " ETP_RULES_INIT_PP,\n", @@ -754,8 +853,8 @@ " batched_tag = \" [batched]\" if p in BATCHED_PRIMITIVES else \"\"\n", " print(f\" {p.name}{batched_tag}\")\n", "\n", - "print(\"\\nTrace propagation rules (ETP_RULES_YW_TO_W):\")\n", - "for p in sorted(ETP_RULES_YW_TO_W.keys(), key=lambda p: p.name):\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", @@ -787,33 +886,156 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 16, "id": "c7d8e9f0", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.939408Z", - "iopub.status.busy": "2026-04-18T06:22:58.939228Z", - "iopub.status.idle": "2026-04-18T06:22:58.943089Z", - "shell.execute_reply": "2026-04-18T06:22:58.942503Z" + "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": [], - "source": "import braintrace\nfrom braintrace import register_primitive\n\n\n# Step 1: Define the implementation.\n# Plain JAX function — no special annotations needed.\ndef _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.\ndef _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\nscaled_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\nprint(\"Primitive registered:\", scaled_mm_p)\nprint(\"Type:\", type(scaled_mm_p).__name__)" + "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": null, + "execution_count": 17, "id": "d8e9f0a1", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.945237Z", - "iopub.status.busy": "2026-04-18T06:22:58.945089Z", - "iopub.status.idle": "2026-04-18T06:22:58.951267Z", - "shell.execute_reply": "2026-04-18T06:22:58.950682Z" + "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": [], - "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\ndef _scaled_yw_to_w(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\ndef _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\ndef _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\ndef _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\nscaled_mm_p.register_etp_rules(\n yw_to_w=_scaled_yw_to_w,\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\nprint(\"yw_to_w registered: \", scaled_mm_p in ETP_RULES_YW_TO_W)\nprint(\"xy_to_dw registered: \", scaled_mm_p in ETP_RULES_XY_TO_DW)\nprint(\"init_drtrl registered:\", scaled_mm_p in ETP_RULES_INIT_DRTRL)\nprint(\"init_pp registered: \", scaled_mm_p in ETP_RULES_INIT_PP)" + "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", @@ -821,10 +1043,10 @@ "id": "e9f0a1b2", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:58.953466Z", - "iopub.status.busy": "2026-04-18T06:22:58.953219Z", - "iopub.status.idle": "2026-04-18T06:22:59.139119Z", - "shell.execute_reply": "2026-04-18T06:22:59.138200Z" + "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": [ @@ -832,13 +1054,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Output shape : (4, 5)\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ + "Output shape : (4, 5)\n", "Matches 2·xw : True\n" ] }, @@ -874,10 +1090,10 @@ "id": "f0a1b2c3", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:59.141017Z", - "iopub.status.busy": "2026-04-18T06:22:59.140842Z", - "iopub.status.idle": "2026-04-18T06:22:59.259288Z", - "shell.execute_reply": "2026-04-18T06:22:59.258589Z" + "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": [ @@ -901,6 +1117,14 @@ "text": [ "Vmap output shape: (8, 4, 5)\n" ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/chaoming/miniconda3/lib/python3.13/site-packages/jax/_src/interpreters/batching.py:603: UserWarning: ETP primitive 'etp_scaled_mm' was decomposed into standard JAX ops under vmap (batch dims (0, None)). Its trainable parameters will NOT be recognized by the eligibility-trace compiler if this trace is compiled for online learning. Map over the data input only (weights unbatched), or call the batched op directly.\n", + " return batcher(vals, dims, **params)\n" + ] } ], "source": [ @@ -924,7 +1148,9 @@ "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." + "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", @@ -964,24 +1190,102 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 20, "id": "5a9cdb2b", "metadata": { "execution": { - "iopub.execute_input": "2026-04-18T06:22:59.270891Z", - "iopub.status.busy": "2026-04-18T06:22:59.270644Z", - "iopub.status.idle": "2026-04-18T06:22:57.323492Z", - "shell.execute_reply": "2026-04-18T06:22:57.322941Z" + "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\nclass 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 * jax.random.normal(jax.random.PRNGKey(0), (hid_dim, hid_dim))\n )\n # Input weight: plain matmul, learned via BPTT instead.\n self.W_in = brainstate.ParamState(\n 0.1 * jax.random.normal(jax.random.PRNGKey(1), (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\nmodel = TinyRNN(in_dim=4, hid_dim=6)\n\n# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.\nalg = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros((2, model.in_dim)), batch_size=2)\n\nprint(\"Compiled ETP relations:\", len(alg.graph.hidden_param_op_relations))\nfor rel in alg.graph.hidden_param_op_relations:\n print(\" primitive =\", rel.primitive.name,\n \" trainable keys =\", list(rel.trainable_vars.keys()))" + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Compiled ETP relations: 1\n", + " primitive = etp_mm trainable keys = ['weight']\n" + ] + } + ], + "source": [ + "import brainstate\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 * jax.random.normal(jax.random.PRNGKey(0), (hid_dim, hid_dim))\n", + " )\n", + " # Input weight: plain matmul, learned via BPTT instead.\n", + " self.W_in = brainstate.ParamState(\n", + " 0.1 * jax.random.normal(jax.random.PRNGKey(1), (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\nETP 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\nWhere 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| `yw_to_w` | $\\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\nFurther reading: `advanced/limitations.ipynb` explains the non-parametric-tail invariant and walks through `GRUCell` (3 Linears, only 2 ETP relations)." + "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)." + ] } ], "metadata": { @@ -1005,4 +1309,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +}