Migrate examples, docs and docstrings onto compile / etrace_grad / etrace_evolve

[chaoming0625]

Why

braintrace.compile has been the documented entry point since before 0.2, and #151 landed the two sequence drivers. Neither migration finished: #151's spec scoped a list of sites, main moved underneath it twice (#148 reorganized the tutorials, #150 rewrote the API reference), and examples/ has no CI coverage — CI.yml runs pytest braintrace/ and nothing else.

So the answer to "how do I train a model with braintrace?" depended on which file you opened:

what the reader found	where
compile(...) + etrace_grad(...)	examples/drtrl/01, docs/advanced/batching.ipynb
compile(...), then a hand-written scan-accumulate	docs/quickstart/quickstart.ipynb, docs/tutorials/drtrl.ipynb
Algo(model) + init_all_states + compile_graph + Vmap	examples/004, examples/100-gru 'batch' arm

Two shipped bugs fixed

Both invisible to CI.

1. examples/drtrl/02-batching-vmap.py raised AttributeError: 'Vmap' object has no attribute 'etrace_grad'. Sequence driver API (etrace_grad / etrace_evolve), and a test suite that finishes #151 migrated its loop to etrace_grad but left the wrapper a bare brainstate.nn.Vmap; only ETraceVmap carries the drivers.
2. examples/004-feedforward-conv-snn.py:302 raised TypeError: ParamDimVjpAlgorithm.__init__() got multiple values for argument 'model' — the model passed both positionally and by keyword, plus a decay_or_rank positional D_RTRL does not take. It survived because OnlineBatchTrainer is never instantiated in main() and 004 is skipped for want of NMNIST.

The compile migration has nowhere to put either bad argument. Since no test reaches 004, both trainers were driven directly on a synthetic batch with tonic stubbed — that run is what proves bug 2 fixed rather than merely rewritten, and it confirms both trainers compile the same trace graph.

All five # kept manual: reasons were wrong

Checked against _compile.py. Three sites migrated. Two keep their code with a corrected reason: drtrl/02 because it is the worked expansion of compile(vmap=True), and pp_prop/05 because it had been routing through compile via _shared since #151 while its note and docstring still described the manual path. The 003-* benchmarks' reasons are restated as the precise limitations — a third state scheme compile does not offer, and a ShapeDtypeStruct example input its vmap axis-strip cannot subscript.

Docs and docstrings

One scan-accumulate block, byte-identical in three notebooks, becomes etrace_grad(reduction='mean'); concepts collapses to reduction='sum'. Four prose passages still named brainstate.transform.scan as the mechanism — that matters as much as the code, since it is why a reader would write a scan after reading migrated code.

Nine docstring Examples blocks built a learner and stopped at one forward call, so the per-class API pages never showed a driver. Two differ on purpose: ThreeFactor rides its per-step modulator as a second sequence, and UORO is the only documented example of window mode.

Equivalence

Both driver substitutions are bit-exact against the current code, which is what licenses re-executing the notebooks and requiring unchanged numbers. drtrl and pp_prop text output is identical; quickstart differs only in repr addresses. concepts cell 13 drifts in the last float digits — but the pre-edit notebook executed on the same machine drifts identically, so it is environmental, and cell 16 (the migrated one) is bit-identical across all three runs.

Verification

check	result
pytest examples/	49 passed, 5 skipped
pytest braintrace/ -n 6	2661 passed
mypy braintrace	clean
docstring snippets	15/15 execute
notebooks	6/6 run
004 arms (skipped by tests)	both run

Spec: docs/specs/2026-07-28-compile-and-driver-migration.md, recording the canonical replacements, the measured facts, what stays manual and why, and four declared limitations found while auditing (ETraceVmap does not forward show_graph/graph/report/param_states; compile(**options) cannot carry model=; no path for the 003-* state scheme; ShapeDtypeStruct fails the vmap axis-strip). Recorded, not fixed — changing compile was out of scope.

Worth flagging: examples/ and the docstring snippets are both documentation that can break silently, and these two bugs are what that costs. braintrace/_docs_examples_test.py was deleted in #150, so nothing runs the snippets now. Wiring either into CI is the change that would stop this recurring.

🤖 Generated with Claude Code

Summary by Sourcery

Migrate examples, tutorials, notebooks and algorithm docstrings onto braintrace.compile plus the etrace_grad/etrace_evolve drivers, fix two broken example trainers, and record the compile/vmap limitations and state schemes in a dedicated spec so the repository consistently teaches and uses a single way to build and drive learners over sequences.

New Features:

• Expose and document etrace_grad and etrace_evolve as the standard sequence drivers for all online-learning algorithms and compile-produced learners.

Bug Fixes:

• Fix examples/drtrl/02-batching-vmap.py by wrapping the learner in braintrace.ETraceVmap so the sequence drivers are available.
• Fix examples/004-feedforward-conv-snn.py batch trainer initialization that previously raised a ParamDimVjpAlgorithm.init argument conflict.

Enhancements:

• Standardize examples and tutorials to construct learners via braintrace.compile rather than manual Algo/init_all_states/compile_graph/vmap wiring.
• Update all relevant notebooks and prose to describe sequence training in terms of etrace_grad and etrace_evolve instead of hand-written scan/for_loop loops.
• Extend algorithm docstrings with end-to-end examples that show both learner construction via compile and sequence driving via etrace_grad, including special cases like ThreeFactor modulators and UORO window mode.
• Document the limitations and state schemes around compile(vmap=True) and ShapeDtypeStruct inputs in a new spec under docs/specs/.

Documentation:

• Refresh quickstart, concepts, drtrl, pp_prop and neural_network_layers notebooks to use etrace_grad/etrace_evolve and align their explanatory text with the compile-and-driver model.
• Add a detailed migration and behavior spec for compile and the sequence drivers in docs/specs/2026-07-28-compile-and-driver-migration.md.

Tests:

• Add and run pytest coverage for examples/, and execute all updated notebooks and algorithm docstring snippets to ensure the migrated APIs and examples behave as documented.

[sourcery-ai]

Sorry @chaoming0625, your pull request is larger than the review limit of 150000 diff characters

[sourcery-ai]

Reviewer's Guide

Migrates examples, tutorials, notebooks, and algorithm docstrings to consistently use braintrace.compile plus the etrace_grad/etrace_evolve sequence drivers, fixes two latent example bugs, and documents compile/driver limitations and the few intentionally manual sites.

Sequence diagram for training and evaluation with compile + etrace drivers

sequenceDiagram
    actor User
    participant Learner as compiled_learner
    participant Model
    participant Optimizer

    %% Training epoch
    User->>Learner: etrace_grad(inputs, targets, step_fn=local_loss, reduction='mean', return_value=true)
    loop for each time_step
        Learner->>Model: local_loss(x_t, y_t)
        Model-->>Learner: step_loss_t, step_grads_t
        Learner->>Learner: accumulate_gradients(step_grads_t)
    end
    Learner-->>User: grads, step_losses
    User->>Optimizer: update(grads)

    %% Evaluation sequence
    User->>Learner: etrace_evolve(inputs, return_outputs=true)
    loop for each time_step
        Learner->>Model: forward(x_t)
        Model-->>Learner: y_t
        Learner->>Learner: update_hidden_and_etrace
    end
    Learner-->>User: predictions

Loading

File-Level Changes

Change	Details	Files
Migrate example trainers from manual algorithm construction and state init to braintrace.compile, including fixing a bad D_RTRL constructor call.	Replace Algo(model) + init_all_states + compile_graph patterns with braintrace.compile(model, Algo, example_inputs[0], batch_size=...) in the GRU and conv-SNN examples. Use braintrace.compile(..., vmap=True) instead of handwritten vmap+compile_graph wiring where appropriate, while keeping one worked manual expansion (drtrl/02-batching-vmap.py) as documentation. Fix examples/004-feedforward-conv-snn.py to construct D_RTRL via compile instead of passing conflicting model and decay_or_rank arguments to ParamDimVjpAlgorithm.	examples/100-gru-on-copying-task.py examples/004-feedforward-conv-snn.py examples/drtrl/02-batching-vmap.py examples/pp_prop/05-batching-vmap.py
Standardize sequence driving on etrace_grad and etrace_evolve instead of handwritten scan or for_loop patterns across notebooks and examples.	In quickstart, drtrl, pp_prop, and concepts notebooks, replace manual scan-accumulate gradient loops and for_loop(learner, inputs) with `learner.etrace_grad(..., reduction='mean'	'sum')andlearner.etrace_evolve(..., return_outputs=True), preserving gradient scaling.</li><li>Update drtrl operator and MNIST examples to use etrace_evolvefor warm-up prefixes while keeping single-step losses for the final timestep.</li><li>Adjust advanced batching tutorial prose to describecompile(..., vmap=True)plusETraceVmap.etrace_grad/etrace_evolve` as the canonical batching/driver path.
Clarify and document cases that intentionally stay manual, especially benchmarking and shape-only compile paths, and correct previous inaccurate "kept manual" comments.	Annotate 003-* memory/speed benchmarks to explain why they keep manual vmap/state schemes and Python loops (measurement of per-step memory) and why those schemes are incompatible with compile’s vmap axis-strip and state init semantics. Update drtrl/02-batching-vmap.py and pp_prop/05-batching-vmap.py comments/docstrings to explain their relationship to compile(..., vmap=True) and drop incorrect "vmap_states not covered" notes. Record declared limitations of compile/ETraceVmap in a new spec, including missing forwarding of introspection methods, inability to pass model= via options, lack of support for certain state schemes, and incompatibility with ShapeDtypeStruct example inputs in vmap mode.	examples/003-snn-memory-and-speed-evaluation-all.py examples/003-snn-memory-and-speed-evaluation-batched.py examples/003-snn-memory-and-speed-evaluation-vmap.py examples/drtrl/02-batching-vmap.py examples/pp_prop/05-batching-vmap.py docs/specs/2026-07-28-compile-and-driver-migration.md
Extend algorithm and compile docstrings to demonstrate sequence drivers and align high-level documentation with the new compile/driver story.	For all main algorithms (D_RTRL, EProp, OSTLRecurrent/Feedforward, pp_prop, ParamDimVjpAlgorithm, IODimVjpAlgorithm, SnAp, ThreeFactor, UORO, OSTL variants), add examples that call learner.etrace_grad (and where relevant, windowed or multi-input usage) after the compile example. Update braintrace.__init__ and _compile.py docs to describe compile as replacing init_all_states+compile_graph(+Vmap) and to mention that compiled learners always carry etrace_grad / etrace_evolve. Add prose to ThreeFactor and UORO docs to explain modulator-as-second-sequence and window mode semantics under etrace_grad.	braintrace/_algorithm/d_rtrl.py braintrace/_algorithm/e_prop.py braintrace/_algorithm/ostl.py braintrace/_algorithm/pp_prop.py braintrace/_algorithm/param_dim_vjp.py braintrace/_algorithm/io_dim_vjp.py braintrace/_algorithm/three_factor.py braintrace/_algorithm/uoro.py braintrace/_algorithm/snap_n.py braintrace/_compile.py braintrace/__init__.py
Fix latent example bugs around missing drivers on vmap wrappers and bad constructor usage, and verify via example tests and snippet execution.	Change examples/drtrl/02-batching-vmap.py to wrap the online model in braintrace.ETraceVmap instead of brainstate.nn.Vmap, fixing AttributeError: 'Vmap' object has no attribute 'etrace_grad'. Ensure vmap wrappers in examples call show_graph() on the underlying learner.module and document that ETraceVmap does not forward introspection methods. Execute pytest examples/ plus direct runs of skipped examples (e.g., 004 with tonic stubbed) and docstring snippets/notebooks to confirm numerical equivalence and that the new drivers work as intended.	examples/drtrl/02-batching-vmap.py examples/004-feedforward-conv-snn.py examples/100-gru-on-copying-task.py braintrace/_algorithm/_sequence_driver_mixin.py (behavior relied on, not modified)

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!