Compile train_synthetic_gradient's window loop instead of tracing it per window

[chaoming0625]

Acts on the one codex finding from the P4 review that I pushed back on and was
wrong about.

My argument was that AGENTS.md rule 10 targets driving the model over time,
while train_synthetic_gradient's per-window loop is a grad-plus-optimiser-step
loop like any training loop. Both halves are true and the conclusion still does
not follow: the body drives the learner — learner(_as_window(...)) under
grad — so a Python loop re-traces the whole custom_vjp machinery every
window. The repo's own quickstart already had the right shape (Python over
epochs, scan over steps, optimiser inside), so the idiom argument pointed this
way too.

Measured with a trace counter on the synthesiser's apply:

windows	Python loop	for_loop
4	14 traces	constant
12	38 traces	constant

Changes

• The window loop is brainstate.transform.for_loop. Everything that varies
  across windows is State and threads through automatically: the model's
  hidden states, the learner's trace, the synthesiser's parameters, and the
  optimiser's moments. Verified against the previous implementation for both the
  built-in SGD step and braintools.optim.Adam — identical per-window errors,
  parameters within float32 rounding (1.5e-8).
• Epochs stay a Python loop: reset calls init_all_states, which reallocates
  state and cannot run under a trace.
• _fit_one returns the traced error rather than calling float() on it. That
  call forced a device sync per window to build a list that was averaged one line
  later; concretisation now happens once, after the loop.
• A ragged final window is refused rather than truncated. for_loop needs
  uniform windows, but the substantive reason is that a short last window fits
  the synthesiser against a shorter future than any window the learner will ever
  be driven with — the same mismatch F-35 documents, introduced by the helper
  instead of by the caller. F-35's entry is narrowed accordingly. No in-tree
  caller was affected; all use divisible lengths.

Tests

• test_the_window_body_is_traced_once_not_once_per_window — compares trace
  counts at two sequence lengths rather than asserting an absolute number.
  Verified to fail on the Python loop (14 vs 38) before being committed.
• test_a_ragged_final_window_is_refused_rather_than_truncated, over three
  (T, chunk) pairs.
• test_a_chunk_size_below_one_is_refused.

Roadmap lesson 48 records the reasoning error: I was arguing about which
category the loop belonged to instead of asking what the body actually did.

Verification

dni_test 30 passed (B4 included), 2526 passed across the rest of the suite,
mypy clean. uoro_test.py was not re-run — dni.py is the only source file
touched and nothing in the UORO path references it.

Summary by Sourcery

Compile train_synthetic_gradient’s per-window loop with a transform-based for_loop, enforce window size consistency, and adjust error handling and documentation around synthetic gradient training.

Bug Fixes:

• Reject ragged final windows so synthesiser training cannot accidentally use a different effective window length than deployment.
• Reject chunk_size values below 1 to prevent invalid window configurations.

Enhancements:

• Run the windowed synthetic gradient training body under a compiled for_loop so it is traced once per epoch instead of once per window.
• Return traced errors from the per-window fit routine and concretise them once at epoch end to avoid per-window device synchronisation.
• Update synthetic gradient training docs and limitations to describe the new tracing behaviour and the narrowed F-35 mismatch contract.

Documentation:

• Clarify train_synthetic_gradient’s input constraints, raised errors, and tracing behaviour in the DNI algorithm docs.
• Record the reasoning mistake and resulting design change in the algorithm-axes roadmap.
• Note that helper-induced window-size mismatches are now prevented in the known limitations spec.

Tests:

• Add tests that verify the window body is traced a constant number of times across different sequence lengths.
• Add tests that enforce rejection of ragged final windows and chunk sizes below one in synthetic gradient training.

[sourcery-ai]

🧙 Sourcery has finished reviewing your pull request!

---

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!