docs: migrate learner workflows to sequence-driver APIs

[poilsosart]

Description

This PR completes the migration of application-facing online-learning workflows to the public learner API:

• construct learners with braintrace.compile(...)
• advance warm-up or evaluation prefixes with learner.etrace_evolve(...)
• accumulate online sequence gradients with learner.etrace_grad(...)
• call vmapped sequence drivers on the returned wrapper rather than learner.module
• preserve existing reset, warm-up, masking, reduction, final-step objective, optimizer parameter, and BPTT semantics
• retain the lower-level paths used by compiler/primitive oracles and memory benchmarks

The Algorithms API documentation, compile docstring, tutorials, notebooks, and affected D-RTRL examples have been updated consistently.

No algorithm implementation or public function signature is changed.

How Has This Been Tested

• python -m pytest braintrace/_compile_test.py braintrace/_algorithm/sequence_test.py -q (134 passed)
• python -m pytest examples/drtrl/tests/test_smoke.py -q (9 passed, 1 skipped)
• Executed docs/advanced/custom_algorithms.ipynb
• Parsed all modified notebooks as JSON
• python -m sphinx -W --keep-going -E -a -b html docs docs/_build/html
• Full python -m pytest -q: 2590 passed, 5 skipped, 4 deselected, and 120 failed in the current environment. The failures are associated with the third-party saiunit.math.exprel(mV) compatibility issue and known order-dependent D-RTRL smoke behavior.

Types of changes

• Documentation (non-breaking change which updates documentation)

Checklist

• Code follows the code style of this project.
• Updated the necessary documentation and examples.
• No generated API RST, CSS, JavaScript, or build artifacts are included.
• All tests pass locally; see the documented third-party and test-order failures above.

Other information

The low-level benchmark and oracle paths remain intentionally unchanged because replacing them would alter the measured workload or correctness oracle.

Summary by Sourcery

Migrate online-learning documentation, tutorials, and D-RTRL examples to the public learner workflow built around braintrace.compile, etrace_evolve, and etrace_grad, clarifying vmap behavior and sequence-driver usage.

Enhancements:

• Refine D-RTRL example scripts to construct learners via braintrace.compile and use sequence-driver APIs for warm-up prefixes and step-wise gradients, including vmap batching and final-step objectives.

Documentation:

• Update the advanced custom algorithm tutorial to use compiled learner objects and sequence drivers instead of direct algorithm construction in runnable examples.
• Revise the quickstart concepts notebook to demonstrate online training with learner.etrace_evolve and learner.etrace_grad rather than raw JAX gradients.
• Adjust RNN and SNN online-learning tutorials to describe warm-up and training phases in terms of learner sequence drivers.
• Clarify the braintrace.compile API docs for vmap mode, emphasizing the ETraceVmap wrapper and how to drive sequences via the learner APIs.

[sourcery-ai]

Reviewer's Guide

Migrates all user-facing online learning docs and examples from ad-hoc algorithm objects to the unified learner workflow built around braintrace.compile and sequence-driver methods, and clarifies compile()’s vmap semantics and return types.

Sequence diagram for the new learner compile and sequence-driver workflow

sequenceDiagram
    actor User
    participant Model
    participant braintrace_compile as braintrace.compile
    participant Learner

    User->>Model: construct_model()
    User->>braintrace_compile: compile(Model, D_RTRL, x0, batch_size, vmap)
    braintrace_compile-->>User: Learner (ETraceAlgorithm or ETraceVmap)

    User->>Learner: etrace_evolve(prefix_inputs)
    Learner->>Model: forward(prefix_inputs)
    Learner-->>User: prefix_outputs

    User->>Learner: etrace_grad(sequence_inputs, targets, step_fn)
    loop over time steps
        Learner->>Learner: step_fn(step_input, step_target)
        Learner->>Model: forward(step_input)
    end
    Learner-->>User: gradients, losses
    User->>Model: apply_gradients(gradients)

Loading

File-Level Changes

Change	Details	Files
Switch advanced custom-algorithm tutorial to construct and drive learners via braintrace.compile and sequence drivers instead of direct algorithm objects.	Update notebook prose to recommend compile() as the runnable entry point while keeping compile_graph discussion as conceptual background. Replace manual algorithm construction and compile_graph calls with braintrace.compile returns and learner.* attribute usage. Use learner.etrace_evolve to run multi-step forward passes and learner.get_etrace_of for trace inspection. Adapt reset and trace-introspection helpers to operate on learners instead of raw algorithms. Demonstrate etrace_grad usage in the custom GradClippedDRTRL smoke test.	docs/advanced/custom_algorithms.ipynb
Refactor DRTRL batching example to the canonical public learner workflow with vmapped learners and sequence drivers.	Replace explicit D_RTRL wrapping, vmap_new_states initialization, and ETraceVmap wiring with a single braintrace.compile(..., vmap=True) call. Drive training with learner.etrace_grad using an internal step_loss that calls the learner once per step. Simplify imports and comments to remove low-level vmap wiring details and emphasize the public workflow. Update plot title to reflect the new public learner workflow focus.	examples/drtrl/02-batching-vmap.py
Clarify compile() API documentation around vmap mode, return types, and sequence drivers.	Clarify that compile may print a report before returning rather than using parenthetical wording. Document that vmap=True returns an ETraceVmap wrapper whose .module is the unbatched learner and which should be driven via its own sequence drivers. Align the Returns section to emphasize ETraceAlgorithm or ETraceVmap and how to access reports and drivers. Tighten docstring examples to emphasize that learners always carry the two sequence drivers.	braintrace/_compile.py
Update quickstart concepts notebook to use learner naming and sequence-driver workflow for online training and gradients.	Rename trainer variable to learner for consistency with the public API. Rework reset and evaluation helpers to drive sequences via learner.etrace_evolve instead of single-step calls. Replace brainstate.transform.grad usage with learner.etrace_grad, including step_fn and has_aux wiring. Adjust explanatory text to distinguish between etrace_evolve and etrace_grad responsibilities.	docs/quickstart/concepts.ipynb
Migrate DRTRL operator and MNIST examples to use learner.etrace_evolve and learner.etrace_grad for warmup prefixes and final-step objectives.	Replace separate grad() calls with learner.etrace_grad using final_step_loss functions that own the single learner call. Drive warm-up prefixes via learner.etrace_evolve while keeping final-step-only objectives intact. Rename online model variables to learner for consistency with the rest of the docs. Ensure optimizer updates now consume gradients produced by the learner’s sequence driver interface.	examples/drtrl/09-classification-mnist.py examples/drtrl/08-operator-conv.py
Align RNN and SNN online-learning tutorials with the new learner sequence-driver workflow.	Update RNN tutorial narrative to describe warm-up and learning phases in terms of learner.etrace_evolve and learner.etrace_grad. Clarify that online learning now uses sequence drivers composed with jit rather than generic grad/vmap/scan. Adjust SNN tutorial to construct learners via braintrace.compile and drive training via learner.etrace_grad instead of brainstate.transform.grad. Refresh guidance on ES-D-RTRL setup to emphasize the learner object returned by compile.	docs/tutorials/rnn_online_learning.ipynb docs/tutorials/snn_online_learning.ipynb docs/apis/algorithms.rst

---

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.

[sourcery-ai]

Hey - I've reviewed your changes and they look great!

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!