[chaoming0625]

Fixes #186.

Summary

• Cause: JaxprEvalForHiddenPerturbation._eval_while (_compiler/hidden_pertubation.py) detached a while only when one of the loop's own outvars was a hidden outvar. In two cases the loop stayed in the perturbed step and jax.vjp raised the reverse-through-while_loop ValueError:

  • the loop result passes through further equations before the write (V <- while(...)(V) + I, the issue's case);
  • the loop is fed only by the input (V <- 0.9 V + while(...)(x) @ w).

• Fix: compile() now finds every top-level while whose outputs reach any hidden outvar, in one backward reachability pass (_while_eqns_reaching), and detaches exactly those. The direct case is a subset, so its behaviour is unchanged. Each detach still emits the WARNING CONTROL_FLOW_OPAQUE_FWD. The hidden-to-hidden Jacobian already uses forward mode for any transition containing a while.

• A loop that reaches no hidden state (e.g. on the readout only) is deliberately left alone. Detaching it would silently zero dL/dh, so it keeps the loud error.

• docs/advanced/limitations.ipynb: the control-flow section was stale. It said cond / scan / while_loop are "NOT supported" and that jnp.where breaks compilation. It is rewritten to match current behaviour:

  • jit inlined;
  • cond if-converted, with the dead-branch NaN caveat;
  • static scans unrolled or descended;
  • weight-free while supported under single-step;
  • a weight inside a while is an error.

  A new executed example (weight inside vs. outside a while) replaces the cond "will not work" example, which does compile today.

• Spec: docs/specs/2026-10-24-issue-186-while-reaching-hidden.md.

Before / after (issue repro)

# main
while_leak ERROR in etrace_grad: ValueError: Reverse-mode differentiation does not work for lax.while_loop ...
# this branch
python     |grad w| = 0.776120
fori       |grad w| = 0.776120
while_leak |grad w| = 1.478937
while_I    |grad w| = 0.776120

while_leak is a different model, so its value differs. It equals its unrolled twin element-wise (see tests).

Tests

• _compiler/hidden_pertubation_test.py:

  • an indirect loop and an input-fed loop are detached, the perturbed VJP is traceable, values are unchanged, and one WARNING is emitted per loop;
  • a readout-only loop is untouched.

  The old test_non_hidden_while_untouched pinned an input-fed loop that does reach the hidden state. That loop made etrace_grad raise on main, so the test is replaced.

• _algorithm/tests/while_support_test.py:

  • D-RTRL on the leak / input cells equals the unrolled twin element-wise, through etrace_grad and through the single-step oracle;
  • the twin equals BPTT (finite-window oracle);
  • a readout-only loop still raises.

• 6 of the new tests fail without the fix.

• Full suite (local CPU): 2921 passed, 0 failed. No new mypy errors compared with main (the typecheck_and_build job fails on main too).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!