Skip to content

Fix stale yw_to_w / element_wise fn references in doc notebooks #136

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chaoming0625 merged 1 commit into from
Jul 9, 2026
Merged

Fix stale yw_to_w / element_wise fn references in doc notebooks #136

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions changelog.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
121 changes: 97 additions & 24 deletions docs/advanced/limitations.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": [],
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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": [],
Expand Down Expand Up @@ -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",
Expand All @@ -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": [
Expand All @@ -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"
]
}
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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": {
Expand All @@ -402,4 +475,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
}
}
Loading