[chaoming0625]

Fixes #185.

Problem

Under vjp_method='single-step' (the default), a step output computed from a value that is not a new hidden state got no online credit along that path, and nothing warned about it. Examples: a spike read from the pre-reset membrane potential, or a readout weight applied inside the step. Two terms were missing:

1. the parameters' same-step gradient along the bypass path;
2. the bypass learning signal dL^t/dh^{t-1} of a previous hidden value, paired with the previous trace e^{t-1}.

Fix

• _compiler/output_bypass.py (new): finds the forward cone of every parameter leaf and hidden invar that reaches a model output or a non-hidden state output without crossing a new hidden value or a stop_gradient. It duplicates that cone on differentiable copies, with every other input detached, and attaches the duplicate as a pure tangent: select_n(is_finite(d), o, o + (d - sg(d))). Forward values are unchanged, and an overflowing output stays inf. A cone that crosses a while or an equation with effects is cut, with a WARNING (OUTPUT_BYPASS_DROPPED). The bypasses that are found get an INFO (OUTPUT_BYPASSES_HIDDEN). With no bypass, the perturbed jaxpr is exactly as before.
• Executor: feeds the bypass parameter values and the previous hidden values to the copies. The VJP then returns exactly the bypass terms.
• Backward:
  • The parameter bypass is added through dl_to_etws_at_t. The single-step guard still rejects any other parameter gradient.
  • The hidden bypass signal goes through _compute_learning_signal and is contracted with e^{t-1} by a new pure hook, _solve_previous_trace_gradients. It is implemented for the ParamDim, IO-dim and random-projection engines. ParamDim uses the previous filtered trace, and IO-dim uses the bias correction one step younger.
• learning_signal='modulatory' replaces dL/dh, so it compiles without the rewrite. Multi-step mode builds no perturbation and is unchanged; it already differentiated both terms.

Before / after (issue repro, etrace_grad, |grad w| summed)

output	main	this PR	BPTT
pre-reset spike, d_rtrl	0.000000	1.079097	1.079097
pre-reset spike, eprop	0.000000	1.079097	1.079097
read from the new hidden value	1.861094	1.861094	1.861094
plain readout tanh(V) @ w_out, w_out	0.0	16.271866	16.271864

Side effects (both now exact)

• A same-step read of a pure hidden copy (h2^t = h1^{t-1}, from KeyError in HiddenGroup.transition when a HiddenState depends only on inputs (no hidden-to-hidden path) #184) is credited as the hidden bypass of h1^{t-1}. D-RTRL now equals BPTT, so the HIDDEN_COPY_SAME_STEP_READ warning added in Fix KeyError for input-only hidden states and dropped pure-copy edges (#184) #190 (unreleased) is removed.
• F-33 (known limitations) is narrowed. Under single-step, a plain readout such as wout in plain_and_etp_rnn now equals BPTT. A plain parameter that only feeds a hidden state (win) is still zero, and the pin and the doc say so.
• A loss that reads a descended scan's stacked ys instead of its carry now gets the carry-read gradient. The limitation pin in scan_descent_support_test.py now asserts that equality.

Edge cases covered by tests

• mixed outputs (no double counting);
• a two-state group with only one bypass;
• an ETP Linear readout (pytree param);
• a parameter applied directly to the output;
• a chain of two ETP ops;
• a non-hidden state output;
• integer outputs and stop_gradient cuts;
• inf outputs;
• a while on the bypass (warning);
• ThreeFactor (no rewrite);
• pp_prop and UORO hooks;
• the ParamDim filter is read but not advanced;
• vmapped lanes vs per-sample BPTT;
• multi-step finite-window oracle vs BPTT.

Spec: docs/specs/2026-10-24-issue-185-output-bypass-gradient.md. Docs: a new "Outputs That Bypass New Hidden States" section in docs/advanced/limitations.ipynb.

[codecov]

Codecov Report

❌ Patch coverage is 97.22222% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
braintrace/_algorithm/vjp_base.py	83.33%	4 Missing ⚠️
braintrace/_compiler/output_bypass.py	98.09%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!