From c6f5a659a611bc0c4f494ef28c6a4a70971d03fc Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 24 Oct 2026 18:20:54 +0800 Subject: [PATCH 1/5] Add the same-step output-bypass gradient under single-step VJP (#185) --- braintrace/_algorithm/io_dim_vjp.py | 29 ++ braintrace/_algorithm/param_dim_vjp.py | 29 ++ .../_algorithm/random_projection_vjp.py | 44 ++- .../_algorithm/tests/output_bypass_test.py | 296 ++++++++++++++++++ .../tests/scan_descent_support_test.py | 21 +- braintrace/_algorithm/vjp_base.py | 122 +++++++- braintrace/_algorithm/vjp_graph_executor.py | 51 ++- braintrace/_compiler/diagnostics.py | 11 + braintrace/_compiler/graph.py | 23 +- braintrace/_compiler/hidden_pertubation.py | 202 +++++++++++- braintrace/_compiler/output_bypass.py | 271 ++++++++++++++++ braintrace/_compiler/output_bypass_test.py | 220 +++++++++++++ braintrace/_compiler/scan_descent.py | 15 +- ...-10-24-issue-185-output-bypass-gradient.md | 114 +++++++ 14 files changed, 1387 insertions(+), 61 deletions(-) create mode 100644 braintrace/_algorithm/tests/output_bypass_test.py create mode 100644 braintrace/_compiler/output_bypass.py create mode 100644 braintrace/_compiler/output_bypass_test.py create mode 100644 docs/specs/2026-10-24-issue-185-output-bypass-gradient.md diff --git a/braintrace/_algorithm/io_dim_vjp.py b/braintrace/_algorithm/io_dim_vjp.py index 202b5107..deb9d640 100644 --- a/braintrace/_algorithm/io_dim_vjp.py +++ b/braintrace/_algorithm/io_dim_vjp.py @@ -1106,6 +1106,35 @@ def _update_etrace_data( return hist_etrace_vals + def _solve_previous_trace_gradients( + self, + running_index: int, + etrace_h2w_at_t_minus_1: Tuple[ + Dict[ETraceX_Key, jax.Array], + Dict[ETraceDF_Key, jax.Array] + ], + dl_to_hidden_groups: Sequence[jax.Array], + weight_vals: Dict[Path, PyTree], + ) -> Dict[Path, PyTree]: + """Contract the previous trace with the hidden output-bypass signal. + + See :meth:`ETraceVjpAlgorithm._solve_previous_trace_gradients`. The + previous trace is one update younger, so its warm-up bias correction + is indexed by ``running_index - 1``. + """ + dG_weights: Dict[Path, Any] = {path: None for path in self.param_states.keys()} + _solve_IO_dim_weight_gradients( + etrace_h2w_at_t_minus_1, + dG_weights, + dl_to_hidden_groups, + self.graph.hidden_param_op_relations, + weight_vals, + running_index - 1, + self.decay_f, + fast_solve=self.fast_solve, + ) + return dG_weights + def _solve_weight_gradients( self, running_index: int, diff --git a/braintrace/_algorithm/param_dim_vjp.py b/braintrace/_algorithm/param_dim_vjp.py index b61105d4..58cb47ab 100644 --- a/braintrace/_algorithm/param_dim_vjp.py +++ b/braintrace/_algorithm/param_dim_vjp.py @@ -1408,6 +1408,35 @@ def _update_etrace_data( return new_etrace + def _solve_previous_trace_gradients( + self, + running_index: int, + etrace_h2w_at_t_minus_1: Dict[ETraceWG_Key, PyTree], + dl_to_hidden_groups: Sequence[jax.Array], + weight_vals: Dict[Path, PyTree], + ) -> Dict[Path, PyTree]: + """Contract the previous trace with the hidden output-bypass signal. + + See :meth:`ETraceVjpAlgorithm._solve_previous_trace_gradients`. With + trace filters active, the previous *filtered* trace (the filter's + current value, not yet advanced this step) is the one paired. + """ + etrace = dict(etrace_h2w_at_t_minus_1) + if self._trace_filters: + for key, flt in self._trace_filters.items(): + if key in etrace: + etrace[key] = flt.value + dG_weights: Dict[Path, Any] = {path: None for path in self.param_states} + _solve_param_dim_weight_gradients( + etrace, + dG_weights, + dl_to_hidden_groups, + self.graph.hidden_param_op_relations, + weight_vals, + fast_solve=self.fast_solve, + ) + return dG_weights + def _solve_weight_gradients( self, running_index: int, diff --git a/braintrace/_algorithm/random_projection_vjp.py b/braintrace/_algorithm/random_projection_vjp.py index 8738c7a5..eb6704f3 100644 --- a/braintrace/_algorithm/random_projection_vjp.py +++ b/braintrace/_algorithm/random_projection_vjp.py @@ -665,6 +665,36 @@ def _update_etrace_data( # the boundary contraction # ------------------------------------------------------------------ # + def _solve_previous_trace_gradients( + self, + running_index: int, + etrace_h2w_at_t_minus_1: Dict[str, Any], + dl_to_hidden_groups: Sequence[jax.Array], + weight_vals: Dict[Path, PyTree], + ) -> Dict[Path, PyTree]: + """Contract the previous trace with the hidden output-bypass signal. + + See :meth:`ETraceVjpAlgorithm._solve_previous_trace_gradients`. + """ + return self._contract_trace(etrace_h2w_at_t_minus_1, dl_to_hidden_groups) + + def _contract_trace( + self, + etrace: Dict[str, Any], + dl_to_hidden_groups: Sequence[jax.Array], + ) -> Dict[Path, Any]: + # ``sum_g (signal_g . s_tilde_g) * theta_tilde_g`` + dG_weights: Dict[Path, Any] = {path: None for path in self.param_states} + for group in self.graph.hidden_groups: + gi = group.index + signal = u.get_mantissa(dl_to_hidden_groups[gi]) + s_tilde = etrace[_S_TILDE][gi] + scale = jnp.sum(signal * s_tilde) + for path in self._etp_paths_of_group(group): + theta = etrace[_THETA_TILDE][(gi, path)] + _update_dict(dG_weights, path, _scale_tree(theta, scale)) + return dG_weights + def _solve_weight_gradients( self, running_index: int, @@ -679,22 +709,14 @@ def _solve_weight_gradients( A scalar per group times a parameter-shaped array — the whole reason the parameter factor carries no hidden axis. """ - dG_weights: Dict[Path, Any] = {path: None for path in self.param_states} - - for group in self.graph.hidden_groups: - gi = group.index - signal = u.get_mantissa(dl_to_hidden_groups[gi]) - s_tilde = etrace_h2w_at_t[_S_TILDE][gi] - scale = jnp.sum(signal * s_tilde) - for path in self._etp_paths_of_group(group): - theta = etrace_h2w_at_t[_THETA_TILDE][(gi, path)] - _update_dict(dG_weights, path, _scale_tree(theta, scale)) + dG_weights = self._contract_trace(etrace_h2w_at_t, dl_to_hidden_groups) # the non-etrace weight gradients (reverse-AD, in-window) for path, dg in dl_to_nonetws_at_t.items(): _update_dict(dG_weights, path, dg) - # the in-window direct term for etrace parameters (multi-step only) + # the in-window direct term for etrace parameters (multi-step), or + # the same-step output-bypass term (single-step, #185) if dl_to_etws_at_t is not None: for path, dg in dl_to_etws_at_t.items(): _update_dict(dG_weights, path, dg, error_when_no_key=True) diff --git a/braintrace/_algorithm/tests/output_bypass_test.py b/braintrace/_algorithm/tests/output_bypass_test.py new file mode 100644 index 00000000..6c4faae7 --- /dev/null +++ b/braintrace/_algorithm/tests/output_bypass_test.py @@ -0,0 +1,296 @@ +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Outputs that bypass every new hidden value train correctly (issue #185). + +Under ``vjp_method='single-step'`` a step output computed from a value that is +not a new hidden state -- a spike read from the pre-reset membrane potential, +a plain readout weight applied inside the step -- used to contribute nothing +to the online gradient, silently. Two terms were missing: the parameters' +same-step bypass gradient, and the bypass learning signal +``dL^t/dh^{t-1}`` paired with the previous trace ``e^{t-1}``. + +The models here are single-layer cells with an element-wise transition, so +D-RTRL and EProp are *exact* on them: per-step ``etrace_grad`` must equal +BPTT element-wise. A mixed output (bypass plus a read of the new hidden +value) pins that neither path is counted twice. The approximate engines +(IO-dim, random projection) are checked through their contraction hooks. +""" + +import warnings + +import brainstate +import braintools +import jax +import jax.numpy as jnp +import numpy.testing as npt +import pytest + +import braintrace + +N_IN, N, T = 3, 4, 10 +_sg = braintools.surrogate.Sigmoid() + + +def _shape(batch_size): + return (N,) if batch_size is None else (batch_size, N) + + +class _LIF(brainstate.nn.Module): + def __init__(self, out='pre_reset'): + super().__init__() + self.out = out + self.w = brainstate.ParamState(brainstate.random.randn(N_IN, N)) + + def init_state(self, batch_size=None, **kwargs): + self.V = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + + def update(self, x): + V_pre = 0.9 * self.V.value + braintrace.matmul(x, self.w.value) + spk = _sg(V_pre - 1.0) + self.V.value = V_pre - jax.lax.stop_gradient(spk) + if self.out == 'pre_reset': + return spk + if self.out == 'mixed': + return _sg(self.V.value - 0.5) + 0.3 * jnp.tanh(V_pre) + return _sg(self.V.value - 0.5) + + +class _CurrentLIF(brainstate.nn.Module): + """Two hidden states in one group; only ``V`` has a hidden bypass.""" + + def __init__(self): + super().__init__() + self.w = brainstate.ParamState(brainstate.random.randn(N_IN, N)) + + def init_state(self, batch_size=None, **kwargs): + self.V = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + self.I = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + + def update(self, x): + self.I.value = 0.8 * self.I.value + braintrace.matmul(x, self.w.value) + V_pre = 0.9 * self.V.value + self.I.value + spk = _sg(V_pre - 1.0) + self.V.value = V_pre - jax.lax.stop_gradient(spk) + return spk + + +class _Readout(brainstate.nn.Module): + def __init__(self): + super().__init__() + self.w = brainstate.ParamState(brainstate.random.randn(N_IN, N)) + self.w_out = brainstate.ParamState(brainstate.random.randn(N, 2)) + + def init_state(self, batch_size=None, **kwargs): + self.V = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + + def update(self, x): + self.V.value = 0.9 * self.V.value + braintrace.matmul(x, self.w.value) + return jnp.tanh(self.V.value) @ self.w_out.value + + +_MODELS = { + 'pre_reset': lambda: _LIF('pre_reset'), + 'mixed': lambda: _LIF('mixed'), + 'hidden': lambda: _LIF('hidden'), + 'current': _CurrentLIF, + 'readout': _Readout, +} + + +def _data(n_out): + brainstate.random.seed(42) + xs = brainstate.random.rand(T, N_IN) * 2.0 + return xs, jnp.full((T, n_out), 0.3) + + +def _loss(y, t): + return jnp.mean((y - t) ** 2) + + +_ALGO_KWARGS = { + 'pp_prop': dict(decay_or_rank=0.9), + 'uoro': dict(vjp_method='single-step'), +} + + +def _online(make, algo, **kwargs): + brainstate.random.seed(0) + net = make() + brainstate.nn.init_all_states(net) + xs, ys = _data(2 if isinstance(net, _Readout) else N) + learner = braintrace.compile(net, algo, xs[0], **{**_ALGO_KWARGS.get(algo, {}), **kwargs}) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + return learner.etrace_grad(xs, ys, step_fn=lambda x, y: _loss(learner(x), y)), learner + + +def _bptt(make): + brainstate.random.seed(0) + net = make() + brainstate.nn.init_all_states(net) + xs, ys = _data(2 if isinstance(net, _Readout) else N) + params = net.states(brainstate.ParamState) + + def total_loss(): + return brainstate.transform.for_loop(lambda x, y: _loss(net(x), y), xs, ys).mean() + + return brainstate.transform.grad(total_loss, params)() + + +@pytest.mark.parametrize('model', list(_MODELS)) +@pytest.mark.parametrize('algo', ['d_rtrl', 'eprop']) +def test_exact_algorithms_match_bptt(model, algo): + online, _ = _online(_MODELS[model], algo) + ref = _bptt(_MODELS[model]) + assert set(online) == set(ref) + for path in ref: + assert float(jnp.abs(ref[path]).sum()) > 1e-6 + npt.assert_allclose(online[path], ref[path], rtol=1e-4, atol=1e-6) + + +def test_issue_values(): + # the numbers in the issue body: 1.079097 (pre-reset), 1.861094 (hidden) + for model, expected in [('pre_reset', 1.079097), ('hidden', 1.861094)]: + online, _ = _online(_MODELS[model], 'd_rtrl') + npt.assert_allclose(float(jnp.abs(online[('w',)]).sum()), expected, rtol=1e-5) + + +@pytest.mark.parametrize('model', ['pre_reset', 'readout']) +@pytest.mark.parametrize('chunk_size', [1, 4]) +def test_multi_step_needs_no_bypass(model, chunk_size): + # Multi-step differentiates the parameters and the hidden values directly, + # so it never had the gap; the finite-window oracle pins that. + from braintrace._testing.oracle import ( + assert_param_gradients_close, bptt_param_gradients, chunked_online_param_gradients, + ) + def factory(): + brainstate.random.seed(0) # the same weights on every build + return _MODELS[model]() + + xs = jax.random.uniform(jax.random.PRNGKey(3), (8, 1, N_IN)) * 2.0 + online = chunked_online_param_gradients( + factory, xs, chunk_size=chunk_size, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + assert_param_gradients_close(online, bptt_param_gradients(factory, xs), atol=1e-5) + + +def test_modulatory_signal_compiles_without_bypass(): + brainstate.random.seed(0) + net = _LIF('pre_reset') + brainstate.nn.init_all_states(net) + xs, _ = _data(N) + algo = braintrace.ThreeFactor(net) + algo.compile_graph(xs[0]) + hp = algo.graph.hidden_perturb + assert list(hp.bypass_param_paths) == [] and list(hp.bypass_hidden_paths) == [] + + +@pytest.mark.parametrize('algo', ['pp_prop', 'uoro', 'd_rtrl', 'eprop']) +def test_bypass_reaches_every_engine(algo): + online, learner = _online(_MODELS['pre_reset'], algo) + hp = learner.graph.hidden_perturb + assert list(hp.bypass_param_paths) == [('w',)] + assert list(hp.bypass_hidden_paths) == [('V',)] + g = online[('w',)] + assert bool(jnp.all(jnp.isfinite(g))) and float(jnp.abs(g).sum()) > 1e-6 + + +def _engine(algo, **kwargs): + brainstate.random.seed(0) + net = _LIF('pre_reset') + brainstate.nn.init_all_states(net) + xs, _ = _data(N) + learner = braintrace.compile(net, algo, xs[0], **{**_ALGO_KWARGS.get(algo, {}), **kwargs}) + learner.init_etrace_state() + # roll a few steps so the traces are non-trivial + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + brainstate.transform.for_loop(learner, xs[:4]) + dl = [jnp.linspace(0.1, 1.0, N).reshape(N, 1)] + weights = {path: st.value for path, st in learner.param_states.items()} + return learner, dl, weights + + +@pytest.mark.parametrize('algo', ['pp_prop', 'uoro']) +def test_previous_trace_hook_is_the_weight_contraction(algo): + learner, dl, weights = _engine(algo) + trace = learner._get_etrace_data() + running_index = 4 + prev = learner._solve_previous_trace_gradients(running_index, trace, dl, weights) + # the contraction of the previous trace, one update younger + main = learner._solve_weight_gradients(running_index - 1, trace, dl, weights, {}, None) + for path in weights: + npt.assert_allclose(prev[path], main[path], rtol=1e-6) + + +def test_previous_trace_hook_reads_the_filtered_trace_without_advancing_it(): + learner, dl, weights = _engine('eprop', kappa_filter_decay=0.9) + assert learner._trace_filters + before = {k: f.value for k, f in learner._trace_filters.items()} + trace = learner._get_etrace_data() + prev = learner._solve_previous_trace_gradients(0, trace, dl, weights) + for k, f in learner._trace_filters.items(): + npt.assert_array_equal(f.value, before[k]) + # identical to contracting the filter's current value directly + filtered = {k: before.get(k, v) for k, v in trace.items()} + learner._trace_filters, saved = {}, learner._trace_filters + try: + ref = learner._solve_previous_trace_gradients(0, filtered, dl, weights) + finally: + learner._trace_filters = saved + for path in weights: + npt.assert_allclose(prev[path], ref[path], rtol=1e-6) + + +def test_vmapped_lanes_match_per_sample_bptt(): + # The OnlineVmapTrainer flow: per-lane states, the graph compiled per lane. + B = 3 + xs = jax.random.uniform(jax.random.PRNGKey(5), (T, B, N_IN)) * 2.0 + + brainstate.random.seed(0) + net = _LIF('pre_reset') + model = braintrace.D_RTRL(net) + + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=B) + def init(): + brainstate.nn.init_all_states(net) + model.compile_graph(xs[0, 0]) + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + init() + vmodel = brainstate.nn.Vmap(model, vmap_states='new') + weights = net.states().subset(brainstate.ParamState) + + def step(prev, x): + g = brainstate.transform.grad(lambda x_: ((vmodel(x_) - 0.3) ** 2).sum(), weights)(x) + return jax.tree.map(lambda a, b: a + b, prev, g), None + + online, _ = brainstate.transform.scan( + step, jax.tree.map(jnp.zeros_like, weights.to_dict_values()), xs) + + ref = None + for b in range(B): + brainstate.random.seed(0) + one = _LIF('pre_reset') + brainstate.nn.init_all_states(one) + g = brainstate.transform.grad( + lambda: brainstate.transform.for_loop( + lambda x: ((one(x) - 0.3) ** 2).sum(), xs[:, b]).sum(), + one.states(brainstate.ParamState))() + ref = g if ref is None else jax.tree.map(lambda a, c: a + c, ref, g) + for path in ref: + npt.assert_allclose(online[path], ref[path], rtol=1e-4, atol=1e-5) diff --git a/braintrace/_algorithm/tests/scan_descent_support_test.py b/braintrace/_algorithm/tests/scan_descent_support_test.py index e7eb615a..68e39d4f 100644 --- a/braintrace/_algorithm/tests/scan_descent_support_test.py +++ b/braintrace/_algorithm/tests/scan_descent_support_test.py @@ -327,22 +327,17 @@ def test_single_step_vjp_descended_equals_unrolled_one_step(self): g_u = self._one_step_grads(UNROLL(), carry_readout=True) assert_param_gradients_close(g_d, g_u, atol=ATOL_EQUIV) - def test_single_step_ys_readout_drops_same_step_signal(self): - """v1 limitation pin: when the loss reads the hidden state through - the descended scan's stacked ys output (``for_loop(...)[-1]``) - instead of the carry, the same-step loss path bypasses the per-step - hidden perturbation (which is added to the scan's *carry* outvar), - so the single-step learning signal — and hence the one-step ETP - gradient — is zero. This parallels the documented Phase-3 - while-hidden same-step limitation. Multi-step vjp is unaffected - (``test_wholeseq_descended_matches_bptt_diagonal_body`` uses the ys - readout). Fix at the model level: return ``self.h.value`` after the - loop.""" + def test_single_step_ys_readout_matches_carry_readout(self): + """Reading the hidden state through the descended scan's stacked ys + output (``for_loop(...)[-1]``) instead of the carry bypasses the + per-step hidden perturbation (which is added to the scan's *carry* + outvar). That used to zero the one-step ETP gradient; the same-step + output-bypass gradient (#185) now carries the path, so both readouts + give the one-step BPTT gradient.""" g_ys = self._one_step_grads(DESCEND(), carry_readout=False) - np.testing.assert_allclose(np.asarray(g_ys[('w',)]), 0.0, atol=0.0) g_carry = self._one_step_grads(DESCEND(), carry_readout=True) assert float(np.linalg.norm(np.asarray(g_carry[('w',)]))) > 1.0 - + assert_param_gradients_close(g_ys, g_carry, atol=ATOL_EQUIV) def test_default_policy_descends_long_scans(): """Phase 4 default: an over-limit ETP scan compiles without any policy.""" diff --git a/braintrace/_algorithm/vjp_base.py b/braintrace/_algorithm/vjp_base.py index 9271935c..9cfc305c 100644 --- a/braintrace/_algorithm/vjp_base.py +++ b/braintrace/_algorithm/vjp_base.py @@ -42,7 +42,7 @@ ) from braintrace._compiler import ControlFlowPolicy, DEFAULT_MAX_JACOBIAN_ELEMENTS from braintrace._op import is_snap_anchored -from ._common import FixedRandomFeedback +from ._common import FixedRandomFeedback, _update_dict from .axes import ETraceConfig from .base import ETraceAlgorithm from .vjp_graph_executor import ETraceVjpGraphExecutor @@ -514,6 +514,9 @@ def __init__( # compiler metadata -- unlike ``sparse_n``, which needs # ``group.snap`` built at compile time. full_jacobian=config.trace_factorization == 'random_projection', + # A modulatory signal *replaces* dL/dh, so there is no reverse-AD + # signal for an output bypass to complete. + include_output_bypass=config.learning_signal != 'modulatory', ) # super initialization @@ -1105,9 +1108,20 @@ def _update_fn_fwd( # describe the rule without controlling it. (self._inject_exit_cotangent(hiddens, aux) if self.config.learning_signal == 'bootstrapped' else None), + # Single-step with a hidden output bypass (#185): the bypass learning + # signal dL^t/dh^{t-1} pairs with the *previous* trace e^{t-1}. + (old_etrace_vals if self._has_hidden_output_bypass() else None), ) return fwd_out, fwd_res + def _has_hidden_output_bypass(self) -> bool: + """Whether a previous hidden value reaches the step output along a + path that bypasses every new hidden value (single-step only).""" + if not self.graph_executor.is_single_step_vjp: + return False + hidden_perturb = self.graph.hidden_perturb + return hidden_perturb is not None and len(hidden_perturb.bypass_hidden_paths) > 0 + def _update_fn_bwd( self, fwd_res: Any, @@ -1136,6 +1150,7 @@ def _update_fn_bwd( args, # original update(*args) tuple, used by _compute_learning_signal aux, # per-call auxiliary data from _get_update_aux, also used by _compute_learning_signal exit_cotangent, # `bootstrapped`: synthetic cotangent at h^exit, or None + prev_etrace_vals, # single-step hidden output bypass: e^{t-1}, or None ) = fwd_res ( @@ -1193,7 +1208,8 @@ def _update_fn_bwd( dg_non_etrace_params, dg_etrace_params, dg_oth_states, - dg_hid_perturb_or_dl2h + dg_hid_perturb_or_dl2h, + dg_hidden_bypass, ) = jax.tree.unflatten(in_tree, cts_out) # @@ -1208,18 +1224,6 @@ def _update_fn_bwd( # compiled program. # if self.graph_executor.is_single_step_vjp: - if len(dg_etrace_params) != 0: - # Under `vjp_method='single-step'` the ETP weights are updated by - # the RTRL recursion, so the transposed jaxpr must not also hand - # back gradients for them. - raise ValueError( - f'Under vjp_method=\'single-step\' the ETP weight gradients ' - f'come from the eligibility-trace recursion, so the backward ' - f'pass must not return any. It returned ' - f'{len(dg_etrace_params)} for ' - f'{sorted(map(str, dg_etrace_params))}. The compiled graph ' - f'and the graph executor disagree; recompile the graph.' - ) hidden_perturb = self.graph.hidden_perturb if hidden_perturb is None: raise ValueError( @@ -1228,6 +1232,22 @@ def _update_fn_bwd( 'graph carries no hidden perturbation. Call `compile_graph()` ' 'on this algorithm before running the backward pass.' ) + bypass_paths = set(hidden_perturb.bypass_param_paths) + if set(dg_etrace_params) != bypass_paths: + # Under `vjp_method='single-step'` the parameters' through-hidden + # credit comes from the eligibility trace; the transposed jaxpr + # may only hand back the same-step *bypass* gradient of the + # parameters the compiler found one for (#185). Anything else + # would count a path twice. + raise ValueError( + f'Under vjp_method=\'single-step\' the parameter gradients ' + f'come from the eligibility-trace recursion, so the backward ' + f'pass may only return the same-step bypass gradients of ' + f'{sorted(map(str, bypass_paths))}. It returned ' + f'{len(dg_etrace_params)} for ' + f'{sorted(map(str, dg_etrace_params))}. The compiled graph ' + f'and the graph executor disagree; recompile the graph.' + ) if len(hidden_perturb.perturb_vars) != len(dg_hid_perturb_or_dl2h): raise ValueError( f'The backward pass returned {len(dg_hid_perturb_or_dl2h)} ' @@ -1283,6 +1303,25 @@ def _update_fn_bwd( dl2h_at_t_or_t_minus_1, (*args, aux) ) + # + # [3a] Single-step hidden output bypass (#185): an output read from + # h^{t-1} without going through h^t carries the learning signal + # dL^t/dh^{t-1}, which pairs with the previous trace e^{t-1}. Solved + # *before* the main contraction, which may advance stateful trace + # filters. + # + dg_weights_prev = None + if prev_etrace_vals is not None: + hidden_perturb = self.graph.hidden_perturb + dl2h_prev = hidden_perturb.bypass_hidden_data_to_hidden_group_data( + [u.get_mantissa(dg_hidden_bypass[path]) for path in hidden_perturb.bypass_hidden_paths], + self.graph.hidden_groups, + ) + dl2h_prev = self._compute_learning_signal(dl2h_prev, (*args, aux)) + dg_weights_prev = self._solve_previous_trace_gradients( + running_index, prev_etrace_vals, dl2h_prev, weight_vals, + ) + # # [3b] `learning_signal='bootstrapped'`: the second linear pass. # @@ -1334,6 +1373,7 @@ def _update_fn_bwd( dg_etrace_params_future, # plain paths only: see below dg_oth_states_future, _dg_hid_perturb_future, # ditto; the signal is pass 1 only + _dg_hidden_bypass_future, # ditto ) = jax.tree.unflatten(in_tree, cts_future) def _add_leaf(x: Any, y: Any) -> Any: @@ -1382,6 +1422,10 @@ def _add(a: Any, b: Any) -> Any: dg_non_etrace_params, dg_etrace_params, ) + if dg_weights_prev is not None: + for path, dg in dg_weights_prev.items(): + if dg is not None: + _update_dict(dg_weights, path, dg) # The returned hidden cotangent carries the future term too: it is what # the *previous* window will receive, and what a synthesiser regresses on. @@ -1756,6 +1800,56 @@ def _project(B: Any, s: Any) -> Any: for gid, s in enumerate(signals) ] + def _solve_previous_trace_gradients( + self, + running_index: int, + etrace_h2w_at_t_minus_1: Any, + dl_to_hidden_groups: Sequence[jax.Array], + weight_vals: Any, + ) -> Dict[Path, Any]: + r"""Contract the previous eligibility trace with a learning signal. + + .. note:: + + This is a protocol method that subclasses implement. It is called + only under ``vjp_method='single-step'`` when the compiled graph + has a hidden output bypass. + + Computes + + .. math:: + + \frac{\partial L^t}{\partial h^{t-1}}\Big|_{\mathrm{bypass}} + \cdot e^{t-1}, + + the credit of an output read from the previous hidden value along a + path that bypasses every new hidden value. The implementation must be + pure: it must not advance any per-step state (trace filters, ...). + + Parameters + ---------- + running_index : int + The running index of the current step (the count of earlier + ``update()`` calls). + etrace_h2w_at_t_minus_1 : Any + The eligibility trace before this step's update. + dl_to_hidden_groups : Sequence[jax.Array] + The bypass learning signal, one array per hidden group. + weight_vals : Any + The weight values. + + Returns + ------- + dict + Parameter path to its gradient (``None`` for no contribution). + """ + raise NotImplementedError( + f'{type(self).__name__} does not implement ' + f'`_solve_previous_trace_gradients`, which the hidden output bypass ' + f'of the compiled graph (hidden states ' + f'{list(self.graph.hidden_perturb.bypass_hidden_paths)}) needs.' + ) + def _solve_weight_gradients( self, running_index: int, diff --git a/braintrace/_algorithm/vjp_graph_executor.py b/braintrace/_algorithm/vjp_graph_executor.py index cc4cae93..40505121 100644 --- a/braintrace/_algorithm/vjp_graph_executor.py +++ b/braintrace/_algorithm/vjp_graph_executor.py @@ -183,6 +183,11 @@ class ETraceVjpGraphExecutor(ETraceGraphExecutor): estimator of an already-block-diagonal recursion would be strictly worse than the recursion itself, so UORO rolls the whole transition. Default ``False``. + include_output_bypass : bool, optional + Under ``vjp_method='single-step'``, add the exact same-step gradient + along paths to the output that bypass every new hidden value: of the + parameters directly, and of the previous hidden values paired with + the previous eligibility trace. Default ``True``. """ __module__ = 'braintrace' @@ -195,6 +200,7 @@ def __init__( snap_max_jacobian_elements: int = DEFAULT_MAX_JACOBIAN_ELEMENTS, control_flow: Optional[ControlFlowPolicy] = None, full_jacobian: bool = False, + include_output_bypass: bool = True, ): super().__init__( model, @@ -213,6 +219,8 @@ def __init__( # Whether ``_compute_hid2hid_jacobian`` keeps the full transition # Jacobian rather than extracting per-position blocks. self.full_jacobian = full_jacobian + # Whether the single-step graph adds the same-step bypass gradients. + self.include_output_bypass = include_output_bypass @property def is_single_step_vjp(self) -> bool: @@ -263,6 +271,7 @@ def compile_graph(self, *args: Any) -> None: sparse_n=self.sparse_n, snap_max_jacobian_elements=self.snap_max_jacobian_elements, control_flow=self.control_flow, + include_output_bypass=self.include_output_bypass, ) def _compute_hid2weight_jacobian( @@ -729,15 +738,28 @@ def solve_h2w_h2h_l2h_jacobian( ) = split_dict_states_v2(self.states) if self.is_single_step_vjp: - etrace_param_vals = dict() assert self.graph.hidden_perturb is not None + # The parameters stay constants of the step: their through-hidden + # credit comes from the eligibility trace. Only the parameters + # whose path to an output bypasses every new hidden value are + # differentiated, and only through their bypass copies (#185). + bypass_paths = tuple(self.graph.hidden_perturb.bypass_param_paths) + etrace_param_vals = {path: etrace_param_states[path].value for path in bypass_paths} hidden_perturbs = self.graph.hidden_perturb.init_perturb_data() etrace_weight_vals_restore = {path: st.value for path, st in etrace_param_states.items()} + # Likewise the previous hidden values whose path to an output + # bypasses every new hidden value: their cotangent is the bypass + # learning signal ``dL^t/dh^{t-1}`` paired with ``e^{t-1}``. + hidden_bypass_vals = { + path: etrace_hidden_states[path].value + for path in self.graph.hidden_perturb.bypass_hidden_paths + } else: etrace_param_vals = {path: st.value for path, st in etrace_param_states.items()} etrace_weight_vals_restore = {k: v for k, v in etrace_param_vals.items()} hidden_perturbs = [] + hidden_bypass_vals = {} non_etrace_param_vals = {path: st.value for path, st in non_etrace_param_states.items()} etrace_state_vals = {path: st.value for path, st in etrace_hidden_states.items()} @@ -749,10 +771,12 @@ def fun_for_vjp( non_etrace_param_vals_: Any, # non-etrace weights etrace_param_vals_: Any, # etrace weights oth_state_vals_: Any, # other states - perturb_vals_: Any # hidden perturbations, useful when computing \partial L / \partial h + perturb_vals_: Any, # hidden perturbations, useful when computing \partial L / \partial h + hidden_bypass_vals_: Any, # copies of the bypass hidden states' previous values ) -> Any: - # assign state values - if len(etrace_param_vals_) > 0: + # assign state values (single-step: the parameter values are fed + # to the bypass copies only, see ``call_hidden_perturb`` below) + if len(etrace_param_vals_) > 0 and not self.is_single_step_vjp: assign_dict_state_values(etrace_param_states, etrace_param_vals_, write=False) assign_dict_state_values(etrace_hidden_states, etrace_hidden_vals_, write=False) assign_dict_state_values(non_etrace_param_states, non_etrace_param_vals_, write=False) @@ -774,6 +798,16 @@ def fun_for_vjp( inputs, perturb_vals_, old_state_vals, + [ + leaf + for path in self.graph.hidden_perturb.bypass_param_paths + for leaf in jax.tree.leaves(etrace_param_vals_[path]) + ], + [ + leaf + for path in self.graph.hidden_perturb.bypass_hidden_paths + for leaf in jax.tree.leaves(hidden_bypass_vals_[path]) + ], ) else: @@ -807,7 +841,8 @@ def scan_over_multiple_steps( non_etrace_weight_vals_ss: Any, # the non-etrace weights etrace_weight_vals_ss: Any, # the etrace weights other_vals_ss: Any, # the initial other states - hidden_perturbs_ss: Any # the hidden perturbations, only used when is_single_step_vjp is True + hidden_perturbs_ss: Any, # the hidden perturbations, only used when is_single_step_vjp is True + hidden_bypass_vals_ss: Any, # the bypass hidden copies, only used when is_single_step_vjp is True ) -> Any: # processing the inputs information @@ -837,6 +872,7 @@ def scan_fn(carray: Any, x_ss: Dict) -> Any: etrace_weight_vals_ss, other_vals_iter, hidden_perturbs_ss, + hidden_bypass_vals_ss, ) if etrace_stepper is None: @@ -885,7 +921,8 @@ def call_over_single_step( non_etrace_weight_vals_ss: Any, # the non-etrace weights etrace_weight_vals_ss: Any, # the etrace weights other_vals_ss: Any, # the initial other states - hidden_perturbs_ss: Any # the hidden perturbations, only used when is_single_step_vjp is True + hidden_perturbs_ss: Any, # the hidden perturbations, only used when is_single_step_vjp is True + hidden_bypass_vals_ss: Any, # the bypass hidden copies, only used when is_single_step_vjp is True ) -> Any: ( out, @@ -900,6 +937,7 @@ def call_over_single_step( etrace_weight_vals_ss, other_vals_ss, hidden_perturbs_ss, + hidden_bypass_vals_ss, ) return ( (out, hidden_vals_iter, other_vals_iter), @@ -929,6 +967,7 @@ def call_over_single_step( etrace_param_vals, # the inputs (single time) other_state_vals, # the inputs (single time) hidden_perturbs, # the inputs (single time) + hidden_bypass_vals, # the inputs (single time) has_aux=True ) diff --git a/braintrace/_compiler/diagnostics.py b/braintrace/_compiler/diagnostics.py index 44046cf1..d52a3149 100644 --- a/braintrace/_compiler/diagnostics.py +++ b/braintrace/_compiler/diagnostics.py @@ -147,6 +147,17 @@ class DiagnosticKind(str, Enum): # online learning signal (WARNING). Reads on the *next* step are exact. HIDDEN_COPY_SAME_STEP_READ = 'hidden_copy_same_step_read' + # Same-step bypass gradients (see _compiler/output_bypass.py) + # + # ``OUTPUT_BYPASSES_HIDDEN``: parameters reach a step output through a + # value that is not a new hidden state; under the single-step VJP their + # exact same-step gradient along that path is added (INFO). + # ``OUTPUT_BYPASS_DROPPED``: part of such a path crosses a ``while`` or an + # equation with effects, which cannot be duplicated; that part of the + # gradient is zero (WARNING). + OUTPUT_BYPASSES_HIDDEN = 'output_bypasses_hidden' + OUTPUT_BYPASS_DROPPED = 'output_bypass_dropped' + # Structural observations (informational / partial) PRIMITIVE_INSIDE_NESTED_JIT = 'primitive_inside_nested_jit' PRIMITIVE_INSIDE_CONTROL_FLOW = 'primitive_inside_control_flow' diff --git a/braintrace/_compiler/graph.py b/braintrace/_compiler/graph.py index 32113991..6b94145c 100644 --- a/braintrace/_compiler/graph.py +++ b/braintrace/_compiler/graph.py @@ -183,6 +183,8 @@ def call_hidden_perturb( args: Inputs, perturb_data: Sequence[jax.Array], old_state_vals: Optional[Sequence[jax.Array]] = None, + bypass_data: Sequence[jax.Array] = (), + bypass_hidden_data: Sequence[jax.Array] = (), ) -> Tuple[Outputs, ETraceVals, StateVals, TempData]: r"""Run the forward pass with additive perturbations injected at the hidden states. @@ -202,6 +204,16 @@ def call_hidden_perturb( old_state_vals : Sequence[jax.Array] or None, optional The state values to run from. When ``None`` (default) the current values of the compiled model states are used. + bypass_data : Sequence[jax.Array], optional + Values for the bypass parameter copies + (:attr:`HiddenPerturbation.bypass_vars`). Differentiating with + respect to them gives the parameters' same-step gradient along + paths that bypass every new hidden value. Defaults to none. + bypass_hidden_data : Sequence[jax.Array], optional + The previous values of the bypass hidden states + (:attr:`HiddenPerturbation.bypass_hidden_vars`). Differentiating + with respect to them gives the learning signal ``dL^t/dh^{t-1}`` + along paths that bypass every new hidden value. Defaults to none. Returns ------- @@ -219,6 +231,8 @@ def call_hidden_perturb( jaxpr_outs = self.hidden_perturb.eval_jaxpr( jax.tree.leaves((args, old_state_vals)), perturb_data, + bypass_data, + bypass_hidden_data, ) return self.module_info._process(*args, jaxpr_outs=jaxpr_outs) @@ -356,6 +370,7 @@ def compile_etrace_graph( sparse_n: Optional[int] = None, snap_max_jacobian_elements: int = DEFAULT_MAX_JACOBIAN_ELEMENTS, control_flow: Optional[ControlFlowPolicy] = None, + include_output_bypass: bool = True, ) -> ETraceGraph: """Construct the eligibility-trace graph for a given model and inputs. @@ -374,6 +389,11 @@ def compile_etrace_graph( The positional arguments required by the model. include_hidden_perturb : bool, optional Whether to include hidden perturbations in the graph. Default ``True``. + include_output_bypass : bool, optional + Whether the hidden perturbation also yields the same-step gradients + along paths to the output that bypass every new hidden value (of the + parameters and of the previous hidden values). Default ``True``. + Ignored when *include_hidden_perturb* is ``False``. include_recurrent_mixing : bool, optional Hidden-group grouping mode for the hidden-to-hidden transition. When ``False`` (default, "without recurrence"), recurrent ETP mixing @@ -609,7 +629,8 @@ def compile_etrace_graph( hidden_perturb = ( add_hidden_perturbation_from_minfo( - minfo, descended_scan_eqn_ids=descended_eqn_ids) + minfo, descended_scan_eqn_ids=descended_eqn_ids, + include_output_bypass=include_output_bypass) if include_hidden_perturb else None ) diff --git a/braintrace/_compiler/hidden_pertubation.py b/braintrace/_compiler/hidden_pertubation.py index 61c1619b..1f6f3644 100644 --- a/braintrace/_compiler/hidden_pertubation.py +++ b/braintrace/_compiler/hidden_pertubation.py @@ -14,10 +14,11 @@ # ============================================================================== -from typing import Dict, FrozenSet, Set, Sequence, NamedTuple, Any +from typing import Dict, FrozenSet, List, Optional, Set, Sequence, NamedTuple, Any import brainstate import jax.core +import jax.numpy as jnp import brainunit as u from braintrace._compatible_imports import ( @@ -51,6 +52,7 @@ extract_module_info, ModuleInfo, ) +from .output_bypass import OutputBypass, add_output_bypass __all__ = [ 'HiddenPerturbation', @@ -116,11 +118,21 @@ class HiddenPerturbation(NamedTuple): perturb_hidden_paths: Sequence[Path] # the hidden state paths that are perturbed perturb_hidden_states: Sequence[brainstate.HiddenState] # the hidden states that are perturbed perturb_jaxpr: ClosedJaxpr # the perturbed jaxpr + # parameters whose path to an output bypasses every new hidden value + # (see ``_compiler/output_bypass.py``), and the copies of their leaves + bypass_param_paths: Sequence[Path] = () + bypass_vars: Sequence[Var] = () + # hidden states whose previous value reaches an output that way, and the + # copies of their previous values + bypass_hidden_paths: Sequence[Path] = () + bypass_hidden_vars: Sequence[Var] = () def eval_jaxpr( self, inputs: Sequence[jax.Array], - perturb_data: Sequence[jax.Array] + perturb_data: Sequence[jax.Array], + bypass_data: Sequence[jax.Array] = (), + bypass_hidden_data: Sequence[jax.Array] = (), ) -> Sequence[jax.Array]: """Evaluate the perturbed jaxpr. @@ -130,16 +142,38 @@ def eval_jaxpr( The flat input values of the original jaxpr. perturb_data : sequence of jax.Array The perturbation values, one per entry of ``perturb_vars``. + bypass_data : sequence of jax.Array, optional + The values of the bypass parameter copies, one per entry of + ``bypass_vars`` (the leaves of each ``bypass_param_paths`` value). + Differentiating the outputs with respect to them gives exactly the + parameters' same-step bypass gradient. Defaults to none. + bypass_hidden_data : sequence of jax.Array, optional + The previous values of the ``bypass_hidden_paths`` hidden states, + one per entry of ``bypass_hidden_vars``. Differentiating the + outputs with respect to them gives exactly the bypass learning + signal ``dL^t/dh^{t-1}``. Defaults to none. Returns ------- sequence of jax.Array The outputs of the perturbed jaxpr. """ + if len(bypass_data) != len(self.bypass_vars): + raise ValueError( + f'Expected {len(self.bypass_vars)} bypass value(s) for the ' + f'parameters {list(self.bypass_param_paths)}, got {len(bypass_data)}.' + ) + if len(bypass_hidden_data) != len(self.bypass_hidden_vars): + raise ValueError( + f'Expected {len(self.bypass_hidden_vars)} bypass value(s) for the ' + f'hidden states {list(self.bypass_hidden_paths)}, ' + f'got {len(bypass_hidden_data)}.' + ) return jax.core.eval_jaxpr( self.perturb_jaxpr.jaxpr, self.perturb_jaxpr.consts, - *(tuple(inputs) + tuple(perturb_data)) + *(tuple(inputs) + tuple(perturb_data) + tuple(bypass_data) + + tuple(bypass_hidden_data)) ) def init_perturb_data(self) -> Sequence[jax.Array]: @@ -153,6 +187,48 @@ def init_perturb_data(self) -> Sequence[jax.Array]: """ return [jax.numpy.zeros(v.aval.shape, dtype=v.aval.dtype) for v in self.perturb_vars] + def bypass_hidden_data_to_hidden_group_data( + self, + bypass_hidden_data: Sequence[jax.Array], + hidden_groups: Sequence[HiddenGroup], + ) -> Sequence[jax.Array]: + """Assemble per-group arrays from per-hidden bypass data. + + Parameters + ---------- + bypass_hidden_data : sequence of jax.Array + One array per entry of ``bypass_hidden_vars`` (typically the + cotangents ``dL^t/dh^{t-1}`` along the bypass paths). + hidden_groups : sequence of HiddenGroup + The hidden groups to assemble, as in + :meth:`perturb_data_to_hidden_group_data`. + + Returns + ------- + sequence of jax.Array + One concatenated array per hidden group; a perturbed hidden state + without a bypass contributes zeros. + + Raises + ------ + ValueError + If ``bypass_hidden_data`` does not have the same length as + ``bypass_hidden_vars``. + """ + bypass_hidden_data = list(bypass_hidden_data) + if len(bypass_hidden_data) != len(self.bypass_hidden_vars): + raise ValueError( + f'Expected {len(self.bypass_hidden_vars)} bypass value(s) for the ' + f'hidden states {list(self.bypass_hidden_paths)}, ' + f'got {len(bypass_hidden_data)}.' + ) + by_path = dict(zip(self.bypass_hidden_paths, bypass_hidden_data)) + per_perturb = [ + by_path[path] if path in by_path else jnp.zeros(v.aval.shape, v.aval.dtype) + for path, v in zip(self.perturb_hidden_paths, self.perturb_vars) + ] + return self.perturb_data_to_hidden_group_data(per_perturb, hidden_groups) + def perturb_data_to_hidden_group_data( self, perturb_data: Sequence[jax.Array], @@ -296,10 +372,18 @@ def __init__( path_to_state: Dict[Path, brainstate.HiddenState], control_flow: ControlFlowPolicy = DEFAULT_CONTROL_FLOW_POLICY, descended_scan_eqn_ids: FrozenSet[int] = frozenset(), + weight_path_to_invars: Optional[Dict[Path, Sequence[Var]]] = None, + bypass_merge_positions: Sequence[int] = (), ): # necessary data structures self.closed_jaxpr = closed_jaxpr + # Same-step bypass gradients (issue #185): the parameter leaves, and + # the outvar positions (model outputs, non-hidden state outputs) where + # a bypass is attached. ``None`` disables the rewrite. + self.weight_path_to_invars = weight_path_to_invars + self.bypass_merge_positions = tuple(bypass_merge_positions) + # Structured scan descent (Phase 4): scan equations rewritten by # ``scan_descent.apply_scan_descent`` (keyed by ``id(eqn)``) bypass # ``check_unsupported_op`` — their weight usage is legal by @@ -323,6 +407,7 @@ def __init__( ) # Keep an ordered version for deterministic iteration in compile() self._hidden_outvars_ordered = hidden_outvars_ordered + self._hidden_outvar_to_invar = hidden_outvar_to_invar self.path_to_state = path_to_state @@ -393,11 +478,24 @@ def compile(self) -> HiddenPerturbation: outvar_subst.get(v, v) if isinstance(v, Var) else v for v in self.closed_jaxpr.jaxpr.outvars ] + + # [same-step bypass gradients] + bypass = self._add_output_bypass( + new_outvars, + cut_vars={outvar_subst.get(v, v) for v in self._hidden_outvars_ordered}, + ) + param_bypass_keys = [k for k in bypass.bypass_keys if k[0] == 'param'] + hidden_bypass_keys = [k for k in bypass.bypass_keys if k[0] == 'hidden'] + n_param_leaves = sum(len(self.weight_path_to_invars[k[1]]) for k in param_bypass_keys) jaxpr = Jaxpr( constvars=list(self.closed_jaxpr.jaxpr.constvars), - invars=list(self.closed_jaxpr.jaxpr.invars) + list(self.perturb_invars.values()), - outvars=new_outvars, - eqns=self.revised_eqns, + invars=( + list(self.closed_jaxpr.jaxpr.invars) + + list(self.perturb_invars.values()) + + list(bypass.bypass_vars) + ), + outvars=bypass.outvars, + eqns=bypass.eqns, effects=self.closed_jaxpr.jaxpr.effects, debug_info=self.closed_jaxpr.jaxpr.debug_info, ) @@ -411,7 +509,11 @@ def compile(self) -> HiddenPerturbation: perturb_vars=brainstate.util.PrettyList(self.perturb_invars.values()), perturb_hidden_paths=brainstate.util.PrettyList(perturb_hidden_paths), perturb_hidden_states=brainstate.util.PrettyList(perturb_hidden_states), - perturb_jaxpr=revised_closed_jaxpr + perturb_jaxpr=revised_closed_jaxpr, + bypass_param_paths=brainstate.util.PrettyList([k[1] for k in param_bypass_keys]), + bypass_vars=brainstate.util.PrettyList(bypass.bypass_vars[:n_param_leaves]), + bypass_hidden_paths=brainstate.util.PrettyList([k[1] for k in hidden_bypass_keys]), + bypass_hidden_vars=brainstate.util.PrettyList(bypass.bypass_vars[n_param_leaves:]), ) # remove the temporal data @@ -420,6 +522,57 @@ def compile(self) -> HiddenPerturbation: self.hidden_jaxvars_to_remove = set() return info + def _add_output_bypass(self, outvars: List, cut_vars: Set[Var]) -> OutputBypass: + """Rewrite the perturbed step to also yield the same-step bypass + gradients of the parameters and of the previous hidden values (see + ``_compiler/output_bypass.py``). + + The sources are keyed ``('param', path)`` (parameters first, so their + copies come first in ``bypass_vars``) and ``('hidden', path)``. + """ + if self.weight_path_to_invars is None: + return OutputBypass(self.revised_eqns, outvars, [], [], []) + sources: Dict = {('param', p): invars for p, invars in self.weight_path_to_invars.items()} + for ov in self._hidden_outvars_ordered: + sources[('hidden', self.outvar_to_hidden_path[ov])] = [self._hidden_outvar_to_invar[ov]] + bypass = add_output_bypass( + self.revised_eqns, + outvars, + sources=sources, + cut_vars=cut_vars, + merge_positions=self.bypass_merge_positions, + ) + bypass_params = [k[1] for k in bypass.bypass_keys if k[0] == 'param'] + bypass_hidden = [k[1] for k in bypass.bypass_keys if k[0] == 'hidden'] + dropped = [k[1] for k in bypass.dropped_keys] + if bypass_params or bypass_hidden: + emit( + kind=DiagnosticKind.OUTPUT_BYPASSES_HIDDEN, + level=DiagnosticLevel.INFO, + message=( + f'Parameters {bypass_params} and previous hidden values ' + f'{bypass_hidden} reach the step output through a value that ' + f'is not a new hidden state; their exact same-step gradient ' + f'along that path is added to the eligibility-trace gradient.' + ), + context={'params': tuple(bypass_params), 'hidden': tuple(bypass_hidden)}, + ) + if dropped: + emit( + kind=DiagnosticKind.OUTPUT_BYPASS_DROPPED, + level=DiagnosticLevel.WARNING, + message=( + f'Parameters or previous hidden values {dropped} reach the ' + f'step output, bypassing every new hidden state, through a ' + f'while loop or an equation with effects. That part of their ' + f'same-step gradient cannot be computed and is zero. Read the ' + f'output from a new hidden state, or move the loop out of the ' + f'output path.' + ), + context={'sources': tuple(dropped)}, + ) + return bypass + def _eval_pjit(self, eqn: JaxprEqn) -> None: """ Evaluating the pjit primitive. @@ -570,6 +723,8 @@ def add_hidden_perturbation_in_jaxpr( path_to_state: Dict[Path, brainstate.HiddenState], control_flow: ControlFlowPolicy = DEFAULT_CONTROL_FLOW_POLICY, descended_scan_eqn_ids: FrozenSet[int] = frozenset(), + weight_path_to_invars: Optional[Dict[Path, Sequence[Var]]] = None, + bypass_merge_positions: Sequence[int] = (), ) -> HiddenPerturbation: """ Adding perturbations to the hidden states in the jaxpr, and replacing the hidden states with the perturbed states. @@ -596,6 +751,14 @@ def add_hidden_perturbation_in_jaxpr( descended_scan_eqn_ids : frozenset of int, optional ``id()`` values of scan equations rewritten by structured scan descent (Phase 4); exempt from the unsupported-op checks. + weight_path_to_invars : dict, optional + Parameter path to its leaf invars. When given, parameters whose path + to an output at ``bypass_merge_positions`` bypasses every new hidden + value get differentiable copies in the perturbed jaxpr (see + :attr:`HiddenPerturbation.bypass_vars`). ``None`` (default) disables + this. + bypass_merge_positions : sequence of int, optional + Positions of the jaxpr outvars where a bypass is attached. Returns ------- @@ -611,12 +774,15 @@ def add_hidden_perturbation_in_jaxpr( path_to_state=path_to_state, control_flow=control_flow, descended_scan_eqn_ids=descended_scan_eqn_ids, + weight_path_to_invars=weight_path_to_invars, + bypass_merge_positions=bypass_merge_positions, ).compile() def add_hidden_perturbation_from_minfo( minfo: ModuleInfo, descended_scan_eqn_ids: FrozenSet[int] = frozenset(), + include_output_bypass: bool = True, ) -> HiddenPerturbation: """Add hidden-state perturbations from a ``ModuleInfo``. @@ -630,6 +796,11 @@ def add_hidden_perturbation_from_minfo( descended_scan_eqn_ids : frozenset of int, default ``frozenset()`` ``id()`` values of scan equations rewritten by structured scan descent (Phase 4); exempt from the unsupported-op checks. + include_output_bypass : bool, default ``True`` + Whether to add the same-step bypass gradients of the parameters and + of the previous hidden values (see + :attr:`HiddenPerturbation.bypass_param_paths`). Disable it for a + learning signal that replaces ``dL/dh`` (``'modulatory'``). Returns ------- @@ -649,9 +820,26 @@ def add_hidden_perturbation_from_minfo( path_to_state=minfo.retrieved_model_states, control_flow=minfo.control_flow, descended_scan_eqn_ids=descended_scan_eqn_ids, + weight_path_to_invars=minfo.weight_path_to_invars if include_output_bypass else None, + bypass_merge_positions=_bypass_merge_positions(minfo), ) +def _bypass_merge_positions(minfo: ModuleInfo) -> List[int]: + """Outvar positions where a parameter's bypass is attached: the model + outputs and the leaves of non-hidden, non-param states (whose incoming + cotangent the backward pass honours). Hidden-state slots and the extra + Jacobian temporaries are excluded.""" + positions = list(range(minfo.num_var_out)) + k = minfo.num_var_out + for st, outvar in zip(minfo.compiled_model_states, minfo.state_tree_outvars): + n = len(jax.tree.leaves(outvar)) + if not isinstance(st, (brainstate.HiddenState, brainstate.ParamState)): + positions.extend(range(k, k + n)) + k += n + return positions + + def add_hidden_perturbation_in_module( model: brainstate.nn.Module, *model_args: Any, diff --git a/braintrace/_compiler/output_bypass.py b/braintrace/_compiler/output_bypass.py new file mode 100644 index 00000000..b42636c3 --- /dev/null +++ b/braintrace/_compiler/output_bypass.py @@ -0,0 +1,271 @@ +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +r"""Same-step gradients along paths to the output that bypass every new +hidden value (issue #185). + +Under ``vjp_method='single-step'`` a parameter's gradient is assembled from +the eligibility trace, ``dL/dh^t \cdot e^t``, where ``dL/dh^t`` is read off the +hidden perturbation. The parameters themselves are constants of the +transposed step. A path from a parameter to a step output that crosses *no* +new hidden value (a pre-reset spike read as the output, a readout weight +applied inside the step, ...) is therefore in neither term, and its +contribution was dropped silently. + +The same holds for the *previous* hidden values: an output computed from +``h^{t-1}`` without going through ``h^t`` (a spike read from the pre-reset +membrane potential ``0.9 V^{t-1} + x @ w``) carries the learning signal +``dL^t/dh^{t-1}`` that RTRL pairs with the previous trace ``e^{t-1}``, and the +hidden perturbation at ``h^t`` does not see it either. + +This module rewrites the perturbed step so that these *bypass* terms are +computed exactly, and computed alone. The *sources* are the parameter leaves +and the hidden invars, each group under a caller-chosen key: + +1. A forward pass finds, for every source, the *cone* of variables it reaches + without crossing a hidden-state slot or a ``stop_gradient``. A source whose + cone reaches a merge slot (a model output or a non-hidden state output) + has a bypass. +2. Each bypass source gets one new invar per leaf (a differentiable copy), + and every equation of the copies' cone is duplicated. A duplicate consumes + the copies (or earlier duplicates) for cone inputs and a ``stop_gradient`` + of the original variable for every other input. The original equations are + untouched. +3. At each merge slot the duplicate ``d`` is attached to the original value + ``o`` as a pure tangent, + ``select_n(is_finite(d), o, o + (d - stop_gradient(d)))``. The forward + value is exactly ``o`` (the guard keeps an overflowing output ``inf`` + rather than ``inf - inf = nan``), and the reverse pass sends the output + cotangent into the copies along the cone only. + +A cone equation that cannot be duplicated is cut, and the sources whose +bypass it carries are reported. That is a ``while`` anywhere inside it (no +reverse-mode rule) or an equation with effects (duplicating it would repeat +the effect). +""" + +from typing import Dict, FrozenSet, Hashable, List, NamedTuple, Sequence, Set, Tuple + +import jax +import jax.numpy as jnp + +from braintrace._compatible_imports import ( + Jaxpr, + JaxprEqn, + Literal, + Var, + is_while_primitive, + new_jaxpr_eqn, + new_var, + stop_gradient_p, +) +from .hidden_group import _transition_contains_while + +__all__ = [ + 'OutputBypass', + 'add_output_bypass', +] + + +class OutputBypass(NamedTuple): + """The result of :func:`add_output_bypass`. + + Attributes + ---------- + eqns : list of JaxprEqn + The equations of the rewritten step: the originals, the cone + duplicates and the merge equations. + outvars : list + The rewritten jaxpr outvars (merge slots point at the merged values). + bypass_vars : list of Var + The new invars, one per leaf of each bypass source, ordered by + ``bypass_keys`` and then by leaf. + bypass_keys : list + The sources with a bypass, in ``sources`` order. + dropped_keys : list + The sources whose bypass was (partly) cut at a ``while`` or an + effectful equation. + """ + + eqns: List[JaxprEqn] + outvars: List + bypass_vars: List[Var] + bypass_keys: List[Hashable] + dropped_keys: List[Hashable] + + +def _is_blocking(eqn: JaxprEqn) -> bool: + # A ``while`` (anywhere inside the equation) has no transpose rule, and + # duplicating an effectful equation would repeat its effect. + if eqn.effects: + return True + return is_while_primitive(eqn) or _transition_contains_while( + Jaxpr(constvars=[], invars=[], outvars=[], eqns=[eqn])) + + +def _is_inexact(v: Var) -> bool: + return jnp.issubdtype(v.aval.dtype, jnp.inexact) + + +def _find_cones( + eqns: Sequence[JaxprEqn], + sources: Dict[Hashable, Sequence[Var]], + cut_vars: Set[Var], +) -> Tuple[Dict[Var, FrozenSet[Hashable]], Dict[Var, FrozenSet[Hashable]]]: + """Forward reachability from the sources. + + Returns ``(reach, blocked)``. ``reach[v]`` holds the sources reaching + ``v`` along duplicable equations; ``blocked[v]`` holds the sources whose + path to ``v`` crosses a blocking equation. + """ + reach: Dict[Var, FrozenSet[Hashable]] = {} + blocked: Dict[Var, FrozenSet[Hashable]] = {} + for key, invars in sources.items(): + for v in invars: + if isinstance(v, Var) and v not in cut_vars: + reach[v] = reach.get(v, frozenset()) | {key} + for eqn in eqns: + if eqn.primitive is stop_gradient_p: + continue + ins = [v for v in eqn.invars if isinstance(v, Var)] + r = frozenset().union(*(reach.get(v, frozenset()) for v in ins)) + b = frozenset().union(*(blocked.get(v, frozenset()) for v in ins)) + if _is_blocking(eqn): + r, b = frozenset(), b | r + if not (r or b): + continue + for ov in eqn.outvars: + if ov in cut_vars: + continue + if r: + reach[ov] = r + if b: + blocked[ov] = b + return reach, blocked + + +def add_output_bypass( + eqns: Sequence[JaxprEqn], + outvars: Sequence, + *, + sources: Dict[Hashable, Sequence[Var]], + cut_vars: Set[Var], + merge_positions: Sequence[int], +) -> OutputBypass: + """Add the exact same-step bypass gradients of the sources to a step. + + Parameters + ---------- + eqns : sequence of JaxprEqn + The equations of the (perturbed) step, in order. + outvars : sequence + The step's jaxpr outvars. + sources : dict + Key to the source's leaf invars (parameter leaves, hidden invars). + cut_vars : set of Var + Variables the cones never cross: the hidden-state slots of the step, + whose same-step credit the eligibility trace already carries. + merge_positions : sequence of int + Positions of ``outvars`` where a bypass is attached (model outputs and + non-hidden state outputs). + + Returns + ------- + OutputBypass + The rewritten step. When no source has a bypass, ``eqns`` and + ``outvars`` are returned unchanged and ``bypass_vars`` is empty. + """ + eqns = list(eqns) + outvars = list(outvars) + reach, blocked = _find_cones(eqns, sources, cut_vars) + + slots = [ + outvars[k] for k in merge_positions + if isinstance(outvars[k], Var) and _is_inexact(outvars[k]) + ] + reaching: Set[Hashable] = set() + dropped: Set[Hashable] = set() + for v in slots: + reaching |= reach.get(v, frozenset()) + dropped |= blocked.get(v, frozenset()) + order = list(sources) + bypass_keys = [k for k in order if k in reaching] + dropped_keys = [k for k in order if k in dropped] + if not bypass_keys: + return OutputBypass(eqns, outvars, [], [], dropped_keys) + + # -- duplicate the cones of the bypass sources' copies -- # + dup: Dict[Var, Var] = {} + bypass_vars: List[Var] = [] + for key in bypass_keys: + for v in sources[key]: + copy = new_var('', jax.core.ShapedArray(v.aval.shape, v.aval.dtype)) + bypass_vars.append(copy) + if v not in cut_vars: + dup[v] = copy + + detached: Dict[Var, Var] = {} + new_eqns: List[JaxprEqn] = [] + + def _detached(v: Var) -> Var: + if v not in detached: + fresh = new_var('', jax.core.ShapedArray(v.aval.shape, v.aval.dtype)) + new_eqns.append(new_jaxpr_eqn([v], [fresh], stop_gradient_p, {}, set())) + detached[v] = fresh + return detached[v] + + for eqn in eqns: + new_eqns.append(eqn) + if eqn.primitive is stop_gradient_p or _is_blocking(eqn): + continue + if not any(isinstance(v, Var) and v in dup for v in eqn.invars): + continue + if all(ov in cut_vars for ov in eqn.outvars): + continue + dup_invars = [ + v if isinstance(v, Literal) else (dup[v] if v in dup else _detached(v)) + for v in eqn.invars + ] + dup_outvars = [ + new_var('', jax.core.ShapedArray(ov.aval.shape, ov.aval.dtype)) + for ov in eqn.outvars + ] + new_eqns.append(eqn.replace(invars=dup_invars, outvars=dup_outvars)) + for ov, dv in zip(eqn.outvars, dup_outvars): + if ov not in cut_vars: + dup[ov] = dv + + # -- attach each duplicate at its merge slots as a pure tangent -- # + merged: Dict[Var, Var] = {} + + def _like(v: Var, dtype=None) -> Var: + return new_var('', jax.core.ShapedArray(v.aval.shape, dtype or v.aval.dtype)) + + for k in merge_positions: + v = outvars[k] + if not (isinstance(v, Var) and v in dup and _is_inexact(v)): + continue + if v not in merged: + d = dup[v] + d_sg, diff, total, finite, out = _like(d), _like(d), _like(v), _like(d, jnp.bool_), _like(v) + new_eqns.append(new_jaxpr_eqn([d], [d_sg], stop_gradient_p, {}, set())) + new_eqns.append(new_jaxpr_eqn([d, d_sg], [diff], jax.lax.sub_p, {}, set())) + new_eqns.append(new_jaxpr_eqn([v, diff], [total], jax.lax.add_p, {}, set())) + new_eqns.append(new_jaxpr_eqn([d], [finite], jax.lax.is_finite_p, {}, set())) + new_eqns.append(new_jaxpr_eqn([finite, v, total], [out], jax.lax.select_n_p, {}, set())) + merged[v] = out + outvars[k] = merged[v] + + return OutputBypass(new_eqns, outvars, bypass_vars, bypass_keys, dropped_keys) diff --git a/braintrace/_compiler/output_bypass_test.py b/braintrace/_compiler/output_bypass_test.py new file mode 100644 index 00000000..fefe84d6 --- /dev/null +++ b/braintrace/_compiler/output_bypass_test.py @@ -0,0 +1,220 @@ +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""The same-step output-bypass rewrite of the perturbed step (issue #185).""" + +import warnings + +import brainstate +import braintools +import jax +import jax.numpy as jnp +import numpy.testing as npt +import pytest + +import braintrace +from braintrace._compiler.diagnostics import DiagnosticKind, DiagnosticLevel +from braintrace._compiler.output_bypass import add_output_bypass + +N_IN, N = 3, 4 +_sg = braintools.surrogate.Sigmoid() + + +class _LIF(brainstate.nn.Module): + """``V_pre = 0.9 V + x @ w``; the output is read from ``V_pre`` or ``V``.""" + + def __init__(self, out='pre_reset'): + super().__init__() + self.out = out + self.w = brainstate.ParamState(brainstate.random.randn(N_IN, N)) + + def init_state(self, *args, **kwargs): + self.V = brainstate.HiddenState(jnp.zeros(N)) + + def update(self, x): + V_pre = 0.9 * self.V.value + braintrace.matmul(x, self.w.value) + spk = _sg(V_pre - 1.0) + self.V.value = V_pre - jax.lax.stop_gradient(spk) + if self.out == 'pre_reset': + return spk + if self.out == 'overflow': + return jnp.exp(1e3 * V_pre) + if self.out == 'while': + return jax.lax.while_loop(lambda c: c[0] < 3, lambda c: (c[0] + 1, c[1] * 0.5), (0, V_pre))[1] + return _sg(self.V.value - 0.5) + + +class _Readout(brainstate.nn.Module): + """A plain readout ``tanh(V) @ w_out`` applied inside the step.""" + + def __init__(self): + super().__init__() + self.w = brainstate.ParamState(brainstate.random.randn(N_IN, N)) + self.w_out = brainstate.ParamState(brainstate.random.randn(N, 2)) + + def init_state(self, *args, **kwargs): + self.V = brainstate.HiddenState(jnp.zeros(N)) + + def update(self, x): + self.V.value = 0.9 * self.V.value + braintrace.matmul(x, self.w.value) + return jnp.tanh(self.V.value) @ self.w_out.value + + +def _graph(model, x=None, **kwargs): + brainstate.nn.init_all_states(model) + x = jnp.linspace(0.5, 1.5, N_IN) if x is None else x + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + return braintrace.compile_etrace_graph(model, x, **kwargs), x + + +def _kinds(graph, level=None): + return [r.kind for r in graph.diagnostics if level is None or r.level == level] + + +class TestDetection: + def test_pre_reset_output_has_param_and_hidden_bypass(self): + brainstate.random.seed(0) + graph, _ = _graph(_LIF('pre_reset')) + hp = graph.hidden_perturb + assert list(hp.bypass_param_paths) == [('w',)] + assert list(hp.bypass_hidden_paths) == [('V',)] + assert len(hp.bypass_vars) == 1 and len(hp.bypass_hidden_vars) == 1 + assert DiagnosticKind.OUTPUT_BYPASSES_HIDDEN in _kinds(graph, DiagnosticLevel.INFO) + + def test_output_from_new_hidden_has_no_bypass(self): + brainstate.random.seed(0) + graph, _ = _graph(_LIF('hidden')) + ref, _ = _graph(_LIF('hidden'), include_output_bypass=False) + hp = graph.hidden_perturb + assert list(hp.bypass_param_paths) == [] and list(hp.bypass_hidden_paths) == [] + # no bypass found -> the perturbed jaxpr is not touched at all + assert len(hp.perturb_jaxpr.jaxpr.eqns) == len(ref.hidden_perturb.perturb_jaxpr.jaxpr.eqns) + assert DiagnosticKind.OUTPUT_BYPASSES_HIDDEN not in _kinds(graph) + + def test_plain_readout_is_param_bypass_only(self): + brainstate.random.seed(0) + graph, _ = _graph(_Readout()) + hp = graph.hidden_perturb + assert list(hp.bypass_param_paths) == [('w_out',)] + assert list(hp.bypass_hidden_paths) == [] + + def test_disabled(self): + brainstate.random.seed(0) + graph, _ = _graph(_LIF('pre_reset'), include_output_bypass=False) + hp = graph.hidden_perturb + assert list(hp.bypass_param_paths) == [] and list(hp.bypass_hidden_paths) == [] + assert len(hp.perturb_jaxpr.jaxpr.invars) == ( + len(graph.module_info.closed_jaxpr.jaxpr.invars) + len(hp.perturb_vars)) + + def test_while_on_bypass_path_is_dropped_with_warning(self): + brainstate.random.seed(0) + graph, _ = _graph(_LIF('while')) + hp = graph.hidden_perturb + assert list(hp.bypass_param_paths) == [] and list(hp.bypass_hidden_paths) == [] + dropped = [r for r in graph.diagnostics if r.kind == DiagnosticKind.OUTPUT_BYPASS_DROPPED] + assert len(dropped) == 1 and dropped[0].level == DiagnosticLevel.WARNING + assert set(dropped[0].context['sources']) == {('w',), ('V',)} + + +class TestEvaluation: + def _eval(self, out, x=None): + brainstate.random.seed(0) + model = _LIF(out) + graph, x = _graph(model, x) + hp = graph.hidden_perturb + inputs = jax.tree.leaves((x, [st.value for st in graph.module_info.compiled_model_states])) + w, V = model.w.value, model.V.value + n_out = graph.module_info.num_var_out + + def run(w_copy, v_copy): + outs = hp.eval_jaxpr(inputs, hp.init_perturb_data(), [w_copy], [v_copy]) + return outs[:n_out] + + return graph, run, w, V, x + + def test_forward_values_unchanged(self): + graph, run, w, V, x = self._eval('pre_reset') + ref = graph.module_info.closed_jaxpr + inputs = jax.tree.leaves((x, [st.value for st in graph.module_info.compiled_model_states])) + expected = jax.core.eval_jaxpr(ref.jaxpr, ref.consts, *inputs)[:graph.module_info.num_var_out] + for a, b in zip(run(w, V), expected): + npt.assert_array_equal(a, b) + + def test_copy_gradients_are_the_bypass_partials(self): + _, run, w, V, x = self._eval('pre_reset') + V = V + 0.3 # a non-trivial previous value + gw, gv = jax.grad(lambda a, b: run(a, b)[0].sum(), argnums=(0, 1))(w, V) + # spk = sg(0.9 V + x @ w - 1): the whole output is a bypass + d = jax.vmap(jax.grad(lambda z: _sg(z - 1.0)))(0.9 * V + x @ w) + npt.assert_allclose(gw, jnp.outer(x, d), rtol=1e-6) + npt.assert_allclose(gv, 0.9 * d, rtol=1e-6) + + def test_overflowing_output_stays_inf(self): + # o + (d - sg(d)) would be inf - inf = nan at an overflowing slot + graph, run, w, V, x = self._eval('overflow', x=jnp.full((N_IN,), 3.0)) + out = run(w, V)[0] + assert bool(jnp.any(jnp.isposinf(out))) and not bool(jnp.any(jnp.isnan(out))) + ref = graph.module_info.closed_jaxpr + inputs = jax.tree.leaves((x, [st.value for st in graph.module_info.compiled_model_states])) + npt.assert_array_equal(out, jax.core.eval_jaxpr(ref.jaxpr, ref.consts, *inputs)[0]) + + def test_wrong_number_of_bypass_values_raises(self): + graph, run, w, V, x = self._eval('pre_reset') + hp = graph.hidden_perturb + inputs = jax.tree.leaves((x, [st.value for st in graph.module_info.compiled_model_states])) + with pytest.raises(ValueError, match='parameters'): + hp.eval_jaxpr(inputs, hp.init_perturb_data(), [], [V]) + with pytest.raises(ValueError, match='hidden states'): + hp.eval_jaxpr(inputs, hp.init_perturb_data(), [w], []) + with pytest.raises(ValueError, match='hidden states'): + hp.bypass_hidden_data_to_hidden_group_data([], graph.hidden_groups) + + +class TestRewrite: + """``add_output_bypass`` on hand-written jaxprs.""" + + @staticmethod + def _rewrite(f, *args, sources, cut=(), merge=None): + closed = jax.make_jaxpr(f)(*args) + jaxpr = closed.jaxpr + src = {k: [jaxpr.invars[i] for i in idx] for k, idx in sources.items()} + cut_vars = {jaxpr.outvars[i] for i in cut} + merge = range(len(jaxpr.outvars)) if merge is None else merge + return jaxpr, add_output_bypass( + jaxpr.eqns, jaxpr.outvars, sources=src, cut_vars=cut_vars, merge_positions=list(merge)) + + def test_stop_gradient_cuts_the_cone(self): + _, res = self._rewrite(lambda w, x: jax.lax.stop_gradient(w * x), 1.0, 2.0, sources={'w': [0]}) + assert res.bypass_keys == [] and res.bypass_vars == [] + + def test_integer_outputs_are_not_merged(self): + _, res = self._rewrite(lambda w: (w > 0).astype(jnp.int32), 1.0, sources={'w': [0]}) + assert res.bypass_keys == [] + + def test_cone_stops_at_cut_vars(self): + # h = w * x is a hidden slot; y = h * 2 reaches the output only through it + def f(w, x): + h = w * x + return h, h * 2.0 + + _, res = self._rewrite(f, 1.0, 2.0, sources={'w': [0]}, cut=(0,), merge=(1,)) + assert res.bypass_keys == [] + + def test_only_listed_positions_are_merged(self): + jaxpr, res = self._rewrite(lambda w, x: (w * x, w + x), 1.0, 2.0, sources={'w': [0]}, merge=(1,)) + assert res.bypass_keys == ['w'] + assert res.outvars[0] is jaxpr.outvars[0] + assert res.outvars[1] is not jaxpr.outvars[1] diff --git a/braintrace/_compiler/scan_descent.py b/braintrace/_compiler/scan_descent.py index 34280913..670d8f88 100644 --- a/braintrace/_compiler/scan_descent.py +++ b/braintrace/_compiler/scan_descent.py @@ -37,15 +37,12 @@ perturbation, seam checks, and learning signals are untouched; only the Jacobian extraction and the trace update see the substep axis. -.. note:: Single-step (perturbation) limitation — the per-step hidden - perturbation is added to the descended scan's *carry* outvar. A loss - that reads the hidden state through the scan's stacked ys instead - (e.g. ``for_loop(...)[-1]``) bypasses the perturbation, so its - same-step reverse credit is dropped from the learning signal (the - one-step ETP gradient is zero). Read the state after the loop - (``self.h.value``) to route the output through the perturbed carry. - Multi-step VJP is unaffected. This parallels the Phase-3 while-hidden - same-step limitation. +.. note:: Single-step (perturbation) — the per-step hidden perturbation is + added to the descended scan's *carry* outvar. A loss that reads the + hidden state through the scan's stacked ys instead (e.g. + ``for_loop(...)[-1]``) bypasses the perturbation; its same-step credit + is recovered by the output-bypass rewrite (``output_bypass.py``, + issue #185), so both readouts give the same gradient. .. note:: Body analysis always runs in the default "without recurrence" grouping mode (``include_recurrent_mixing=False``), independent of the diff --git a/docs/specs/2026-10-24-issue-185-output-bypass-gradient.md b/docs/specs/2026-10-24-issue-185-output-bypass-gradient.md new file mode 100644 index 00000000..499de216 --- /dev/null +++ b/docs/specs/2026-10-24-issue-185-output-bypass-gradient.md @@ -0,0 +1,114 @@ +# Issue #185: same-step gradient of parameters whose path to the output bypasses new hidden values + +Issue: [#185](https://github.com/chaobrain/braintrace/issues/185) + +## Problem + +Under `vjp_method='single-step'` (the default), the gradient of every +`ParamState` is assembled from two sources: + +- the eligibility trace: `dL/dh^t · e^t`, where `dL/dh^t` is read off the + hidden perturbation `h^t = fresh + p` (`_compiler/hidden_pertubation.py`); +- the transposed step jaxpr (`dg_etrace_params` in + `VjpAlgorithm._update_fn_bwd`), which **must be empty**. The parameters are + closed over as constants in `ETraceVjpGraphExecutor.solve_h2w_h2h_l2h_jacobian` + (`etrace_param_vals = {}`), because `split_dict_states_v2` classes *every* + `ParamState` as an eligibility-trace parameter. + +A parameter can also reach the step output **without passing through any new +hidden value**. Its same-step contribution then appears in neither source and +is dropped silently: no error, no warning. Measured on `main` (D-RTRL, +`etrace_grad`, |grad| summed): + +| model | online | BPTT | +|---|---|---| +| issue LIF, output `sg(V_pre - th)` (ETP weight through a pre-reset value) | 0.000000 | 1.079097 | +| plain readout inside the step, `tanh(V) @ w_out` (`w_out`) | 0.0 | 16.27 | + +The second row is the same defect, and probably the more common one: any +readout weight applied inside the compiled step gets no gradient. + +## Fix: add the exact same-step bypass term + +For a parameter `w` the true gradient decomposes, per step, into paths that +cross a new hidden value (handled by the trace, unchanged) and paths that do +not, the *bypass*. The bypass contribution is a plain same-step reverse-mode +term with no temporal component, so it is exact for every algorithm, exact or +approximate, and is added the same way for all of them. + +1. **Cone detection (compile time, perturbation pass).** Forward reachability + from every parameter leaf invar through the perturbed step. Propagation stops + at hidden-state output slots and at `stop_gradient`. A parameter whose cone + reaches a model output, or a non-hidden, non-param state output, is a + *bypass parameter*. An INFO diagnostic `OUTPUT_BYPASSES_HIDDEN` names them. +2. **Cone duplication (compile time).** For each bypass parameter the perturbed + jaxpr gets one extra invar per leaf (a differentiable copy of the value). + Every equation in the copies' cone is duplicated, consuming the copies + (or earlier duplicates) for cone inputs and `stop_gradient` of the original + variable for every other input. The original equations are kept unchanged, + so values, the hidden perturbation and every other gradient are untouched. + At each merge slot (model outputs and non-hidden state outputs) the + duplicate `d` is attached as a pure tangent: + `out = select_n(is_finite(d), out, out + (d - stop_gradient(d)))`. + The forward value is bitwise `out`, and the `select_n` guard keeps an + overflowing output from turning into NaN. Hidden slots and the extra + Jacobian temporaries are never merged. +3. **Executor.** In single-step mode, `etrace_param_vals` holds the bypass + parameters (instead of `{}`) and is fed to the copies only. The real + parameter values still come from the states, as constants. The VJP + therefore returns exactly the bypass gradient in `dg_etrace_params`. +4. **Backward.** `_update_fn_bwd` accepts `dg_etrace_params` whose keys are the + compiled bypass parameters (anything else still raises) and passes them as + `dl_to_etws_at_t`. Every `_solve_weight_gradients` already adds this term + (`param_dim_vjp`, `io_dim_vjp`, `random_projection_vjp`). For parameters + with no ETP relation it becomes their only gradient. Bootstrapped DNI is + unaffected: the future cotangent enters at hidden exits, which the cone + never crosses. +5. **Unfixable remainder: WARNING.** A cone equation that cannot be + duplicated is cut, with a WARNING (`OUTPUT_BYPASS_DROPPED`) naming the + parameters. That covers a `while` (not reverse-differentiable) or an + equation with effects (duplicating it would repeat the effect). + +No change when there is no bypass: the perturbed jaxpr and the executor are +exactly as before. Multi-step mode differentiates the parameters directly +and already includes the bypass. A test checks that, and that path is not +changed. + +## Edge cases + +- Output reads the new hidden value: no bypass, jaxpr unchanged. +- Mixed output (bypass + hidden path in one output): both terms, no double + counting (the cone is cut at hidden slots). +- Weight feeds another ETP primitive on the readout path (`W2(W1 x)` with no + hidden state): the cone runs through `W2`'s primitive. The ETP JVP rules are + used, and both get bypass gradients. +- Output is the parameter itself (`return w`). +- Pytree / multi-leaf parameter (weight + bias). +- Integer or boolean outputs: not merged. +- Output overflow (`inf`): the forward value stays `inf`, not NaN. +- `while` on the bypass path: dropped with a WARNING. +- Other (non-hidden) state written from a weight (`ShortTermState`): merged + slot, its incoming cotangent is honoured. +- vmap / batched learners. + +## Tests + +- `_compiler/hidden_pertubation_test.py`: + - detection: bypass and no-bypass; + - structure: no extra invars without a bypass; + - forward values bitwise unchanged; + - the gradient w.r.t. the copies equals a hand-written bypass VJP; + - `while` on the cone gives a WARNING; + - `inf` output stays `inf`. +- `_algorithm/d_rtrl_test.py`, `_algorithm/e_prop_test.py`, `pp_prop_test.py`: + - the issue's LIF and the plain-readout model: D-RTRL equals BPTT, + through `etrace_grad` (issue repro path) and through the finite-window + oracle with chunk size 1; + - mixed output: no double count; + - EProp / pp_prop: the bypass term is present and the algorithms keep + their own trace terms. Checked against the same algorithm on the + `new hidden value` twin model, via the finite-window oracle; + - multi-step mode already equals BPTT on these models (regression pin); + - batched (`batch_size=2`) learner. +- `_algorithm/vjp_base_test.py`: the guard still raises for unexpected ETP + gradients under single-step. From 4b940b4c8e01f2ab3783c9d25482f1164ff55279 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 24 Oct 2026 18:25:53 +0800 Subject: [PATCH 2/5] Credit the previous hidden value's bypass, gate modulatory, tests and spec (#185) --- .../_algorithm/tests/output_bypass_test.py | 93 ++++++-- braintrace/_compiler/diagnostics.py | 10 - braintrace/_compiler/hidden_group.py | 35 --- braintrace/_compiler/hidden_group_test.py | 39 ---- ...-10-24-issue-185-output-bypass-gradient.md | 217 +++++++++++------- 5 files changed, 213 insertions(+), 181 deletions(-) diff --git a/braintrace/_algorithm/tests/output_bypass_test.py b/braintrace/_algorithm/tests/output_bypass_test.py index 6c4faae7..ad3de57c 100644 --- a/braintrace/_algorithm/tests/output_bypass_test.py +++ b/braintrace/_algorithm/tests/output_bypass_test.py @@ -87,6 +87,31 @@ def update(self, x): return spk +class _CopyRead(brainstate.nn.Module): + """A pure copy ``h2^t = h1^{t-1}`` read in the same step (#184). + + In the jaxpr ``h2^t`` and ``h1^{t-1}`` are one variable; the read is + credited as the hidden bypass of ``h1^{t-1}``, which is the same value + with the same trace (``e2^t = e1^{t-1}``). + """ + + def __init__(self): + super().__init__() + self.w = brainstate.ParamState(brainstate.random.randn(N_IN, N)) + + def init_state(self, batch_size=None, **kwargs): + self.h1 = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + self.h2 = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + self.V = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + + def update(self, x): + old1 = self.h1.value + self.V.value = 0.9 * self.V.value + self.h2.value + self.h1.value = 0.5 * old1 + braintrace.matmul(x, self.w.value) + self.h2.value = old1 + return jnp.tanh(self.h2.value) + jnp.tanh(self.V.value) + + class _Readout(brainstate.nn.Module): def __init__(self): super().__init__() @@ -101,14 +126,52 @@ def update(self, x): return jnp.tanh(self.V.value) @ self.w_out.value +class _ReadoutVariants(brainstate.nn.Module): + """Bypass shapes: an ETP ``Linear`` readout (a weight + bias pytree), a + parameter applied directly to the output, a chain of two ETP ops with no + hidden state between them, and a parameter written into a non-hidden + state that is the output.""" + + def __init__(self, kind): + super().__init__() + self.kind = kind + self.w = brainstate.ParamState(brainstate.random.randn(N_IN, N)) + self.out = braintrace.nn.Linear(N, 2) + self.b = brainstate.ParamState(jnp.full((N,), 0.1)) + + def init_state(self, batch_size=None, **kwargs): + self.V = brainstate.HiddenState(jnp.zeros(_shape(batch_size))) + self.s = brainstate.ShortTermState(jnp.zeros(_shape(batch_size))) + + def update(self, x): + self.V.value = 0.9 * self.V.value + braintrace.matmul(x, self.w.value) + if self.kind == 'linear': + return self.out(jnp.tanh(self.V.value)) + if self.kind == 'param': + return jnp.tanh(self.V.value) * self.b.value + self.b.value + if self.kind == 'chain': + return self.out(jnp.tanh(braintrace.matmul(x, self.w.value))) + self.out(self.V.value) + # 'state': no recursion through ``s`` -- a non-hidden state carries no + # trace, so only its same-step credit is online + self.s.value = jnp.tanh(self.V.value) * self.b.value + return self.s.value + + _MODELS = { 'pre_reset': lambda: _LIF('pre_reset'), 'mixed': lambda: _LIF('mixed'), 'hidden': lambda: _LIF('hidden'), 'current': _CurrentLIF, + 'copy_read': _CopyRead, 'readout': _Readout, + 'linear': lambda: _ReadoutVariants('linear'), + 'param': lambda: _ReadoutVariants('param'), + 'chain': lambda: _ReadoutVariants('chain'), + 'state': lambda: _ReadoutVariants('state'), } +_N_OUT = {'readout': 2, 'linear': 2, 'chain': 2} + def _data(n_out): brainstate.random.seed(42) @@ -126,45 +189,47 @@ def _loss(y, t): } -def _online(make, algo, **kwargs): +def _online(model, algo, **kwargs): brainstate.random.seed(0) - net = make() + net = _MODELS[model]() brainstate.nn.init_all_states(net) - xs, ys = _data(2 if isinstance(net, _Readout) else N) + xs, ys = _data(_N_OUT.get(model, N)) learner = braintrace.compile(net, algo, xs[0], **{**_ALGO_KWARGS.get(algo, {}), **kwargs}) with warnings.catch_warnings(): warnings.simplefilter('ignore') return learner.etrace_grad(xs, ys, step_fn=lambda x, y: _loss(learner(x), y)), learner -def _bptt(make): +def _bptt(model): brainstate.random.seed(0) - net = make() + net = _MODELS[model]() brainstate.nn.init_all_states(net) - xs, ys = _data(2 if isinstance(net, _Readout) else N) + xs, ys = _data(_N_OUT.get(model, N)) params = net.states(brainstate.ParamState) def total_loss(): return brainstate.transform.for_loop(lambda x, y: _loss(net(x), y), xs, ys).mean() - return brainstate.transform.grad(total_loss, params)() + # some variants leave a parameter unused + return brainstate.transform.grad(total_loss, params, check_states=False)() @pytest.mark.parametrize('model', list(_MODELS)) @pytest.mark.parametrize('algo', ['d_rtrl', 'eprop']) def test_exact_algorithms_match_bptt(model, algo): - online, _ = _online(_MODELS[model], algo) - ref = _bptt(_MODELS[model]) - assert set(online) == set(ref) + online, _ = _online(model, algo) + ref = _bptt(model) for path in ref: - assert float(jnp.abs(ref[path]).sum()) > 1e-6 - npt.assert_allclose(online[path], ref[path], rtol=1e-4, atol=1e-6) + if all(float(jnp.abs(b).sum()) == 0.0 for b in jax.tree.leaves(ref[path])): + continue # an unused parameter of a variant + for a, b in zip(jax.tree.leaves(online[path]), jax.tree.leaves(ref[path])): + npt.assert_allclose(a, b, rtol=1e-4, atol=1e-5) def test_issue_values(): # the numbers in the issue body: 1.079097 (pre-reset), 1.861094 (hidden) for model, expected in [('pre_reset', 1.079097), ('hidden', 1.861094)]: - online, _ = _online(_MODELS[model], 'd_rtrl') + online, _ = _online(model, 'd_rtrl') npt.assert_allclose(float(jnp.abs(online[('w',)]).sum()), expected, rtol=1e-5) @@ -200,7 +265,7 @@ def test_modulatory_signal_compiles_without_bypass(): @pytest.mark.parametrize('algo', ['pp_prop', 'uoro', 'd_rtrl', 'eprop']) def test_bypass_reaches_every_engine(algo): - online, learner = _online(_MODELS['pre_reset'], algo) + online, learner = _online('pre_reset', algo) hp = learner.graph.hidden_perturb assert list(hp.bypass_param_paths) == [('w',)] assert list(hp.bypass_hidden_paths) == [('V',)] diff --git a/braintrace/_compiler/diagnostics.py b/braintrace/_compiler/diagnostics.py index d52a3149..ff483a7b 100644 --- a/braintrace/_compiler/diagnostics.py +++ b/braintrace/_compiler/diagnostics.py @@ -137,16 +137,6 @@ class DiagnosticKind(str, Enum): CONTROL_FLOW_OPAQUE_FWD = 'control_flow_opaque_fwd' CONTROL_FLOW_RECURRENT_MIXING = 'control_flow_recurrent_mixing' - # Hidden-state copies (see _compiler/hidden_group.py) - # - # ``HIDDEN_COPY_SAME_STEP_READ``: a hidden state is written as a pure copy - # of another hidden state's previous value (``h2^t = h1^{t-1}``, a shift - # register), and that value also reaches a model output or a non-hidden - # state in the same step. In the jaxpr ``h2^t`` and ``h1^{t-1}`` are one - # variable, so such a read cannot be attributed to ``h2^t``; it carries no - # online learning signal (WARNING). Reads on the *next* step are exact. - HIDDEN_COPY_SAME_STEP_READ = 'hidden_copy_same_step_read' - # Same-step bypass gradients (see _compiler/output_bypass.py) # # ``OUTPUT_BYPASSES_HIDDEN``: parameters reach a step output through a diff --git a/braintrace/_compiler/hidden_group.py b/braintrace/_compiler/hidden_group.py index 725514a8..f182a7e2 100644 --- a/braintrace/_compiler/hidden_group.py +++ b/braintrace/_compiler/hidden_group.py @@ -1665,40 +1665,6 @@ def _attach_snap_pattern(self, group: HiddenGroup) -> HiddenGroup: ) return group._replace(snap=pattern) - def _warn_copy_read_in_same_step(self, copy_var: Var) -> None: - """Warn when a pure hidden copy is also read towards an output this step. - - ``copy_var`` is both ``h2^t`` (a pure copy) and ``h1^{t-1}``. A read of - it that reaches a model output or a non-hidden state, without first - passing through a new hidden value, cannot be told apart from a read - of ``h1^{t-1}``. It therefore gets no online learning signal. - """ - reached = {copy_var} - for eqn in self.jaxpr.eqns: - if any(isinstance(v, Var) and v in reached for v in eqn.invars): - # stop at new hidden values: their own perturbation carries the signal - reached.update(ov for ov in eqn.outvars if ov not in self.hidden_outvar_to_invar) - # ``copy_var`` occupies its own hidden-state slot once; any further - # occurrence among the jaxpr outvars is a same-step read. - hits = [v for v in self.jaxpr.outvars if isinstance(v, Var) and v in reached] - if len(hits) <= 1: - return - copy_path = self.outvar_to_hidden_path[copy_var] - source_path = self.invar_to_hidden_path[copy_var] - emit( - kind=DiagnosticKind.HIDDEN_COPY_SAME_STEP_READ, - level=DiagnosticLevel.WARNING, - message=( - f'Hidden state {copy_path} is written as a pure copy of the ' - f'previous value of {source_path}, and that value also reaches ' - f'an output in the same step. In the traced program the two are ' - f'one variable, so this read gets no online learning signal. ' - f'Read {copy_path} on the next step, or compute the output from ' - f'a new hidden value.' - ), - hidden_paths=(copy_path, source_path), - ) - def _post_check(self) -> Tuple[ Sequence[HiddenGroup], Dict[Path, HiddenGroup], @@ -1751,7 +1717,6 @@ def _post_check(self) -> Tuple[ if self.outvar_to_hidden_path[outvar] in self.descended_hidden_paths: continue outvar_groups.append([self.hidden_invar_to_outvar[outvar], outvar]) - self._warn_copy_read_in_same_step(outvar) outvar_groups = _merge_groups_ordered(outvar_groups) outvar_order = {ov: i for i, ov in enumerate(self.hidden_outvar_to_invar)} outvar_groups = [ diff --git a/braintrace/_compiler/hidden_group_test.py b/braintrace/_compiler/hidden_group_test.py index 0a8bc784..18336093 100644 --- a/braintrace/_compiler/hidden_group_test.py +++ b/braintrace/_compiler/hidden_group_test.py @@ -2131,42 +2131,3 @@ def update(self, x): i1, i2 = group.hidden_paths.index(('h1',)), group.hidden_paths.index(('h2',)) np.testing.assert_allclose(jac[..., i2, i1], 1.0) np.testing.assert_allclose(jac[..., i1, i1], 0.5) - - -class TestHiddenCopySameStepRead: - """A same-step read of a pure hidden copy warns; a next-step read does not (#184).""" - - def _kinds(self, read_now): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState(brainstate.random.randn(3, 4)) - - def init_state(self, *args, **kwargs): - self.h1 = brainstate.HiddenState(jnp.zeros(4)) - self.h2 = brainstate.HiddenState(jnp.zeros(4)) - self.V = brainstate.HiddenState(jnp.zeros(4)) - - def update(self, x): - old1 = self.h1.value - self.V.value = 0.9 * self.V.value + self.h2.value - self.h1.value = 0.5 * old1 + braintrace.matmul(x, self.w.value) - self.h2.value = old1 - return jnp.tanh(self.h2.value) if read_now else jnp.tanh(self.V.value) - - net = Net() - brainstate.nn.init_all_states(net) - with diagnostic_context() as reporter, warnings.catch_warnings(): - warnings.simplefilter('ignore') - _, path_to_group = find_hidden_groups_from_module(net, jnp.ones(3)) - assert path_to_group[('h1',)] is path_to_group[('h2',)] - return [r for r in reporter.records() if r.kind == DiagnosticKind.HIDDEN_COPY_SAME_STEP_READ] - - def test_same_step_read_warns(self): - recs = self._kinds(read_now=True) - assert len(recs) == 1 - assert recs[0].level == DiagnosticLevel.WARNING - assert recs[0].hidden_paths == (('h2',), ('h1',)) - - def test_next_step_read_is_silent(self): - assert self._kinds(read_now=False) == [] diff --git a/docs/specs/2026-10-24-issue-185-output-bypass-gradient.md b/docs/specs/2026-10-24-issue-185-output-bypass-gradient.md index 499de216..aa40e676 100644 --- a/docs/specs/2026-10-24-issue-185-output-bypass-gradient.md +++ b/docs/specs/2026-10-24-issue-185-output-bypass-gradient.md @@ -1,4 +1,4 @@ -# Issue #185: same-step gradient of parameters whose path to the output bypasses new hidden values +# Issue #185: same-step gradient along paths to the output that bypass new hidden values Issue: [#185](https://github.com/chaobrain/braintrace/issues/185) @@ -10,105 +10,156 @@ Under `vjp_method='single-step'` (the default), the gradient of every - the eligibility trace: `dL/dh^t · e^t`, where `dL/dh^t` is read off the hidden perturbation `h^t = fresh + p` (`_compiler/hidden_pertubation.py`); - the transposed step jaxpr (`dg_etrace_params` in - `VjpAlgorithm._update_fn_bwd`), which **must be empty**. The parameters are - closed over as constants in `ETraceVjpGraphExecutor.solve_h2w_h2h_l2h_jacobian` + `VjpAlgorithm._update_fn_bwd`), which **had to be empty**. The parameters + are closed over as constants in + `ETraceVjpGraphExecutor.solve_h2w_h2h_l2h_jacobian` (`etrace_param_vals = {}`), because `split_dict_states_v2` classes *every* `ParamState` as an eligibility-trace parameter. -A parameter can also reach the step output **without passing through any new -hidden value**. Its same-step contribution then appears in neither source and -is dropped silently: no error, no warning. Measured on `main` (D-RTRL, -`etrace_grad`, |grad| summed): +A value can reach the step output **without passing through any new hidden +value**, and then its contribution appears in neither source. It is dropped +silently: no error, no warning. Two terms go missing: + +1. **Parameter bypass**: `∂L^t/∂w` along such a path (a plain readout weight + applied inside the step, a weight feeding a pre-reset value). +2. **Hidden bypass**: `∂L^t/∂h^{t-1}` along such a path, which RTRL pairs + with the previous trace `e^{t-1} = ∂h^{t-1}/∂w`. In the issue's LIF, + `spk = sg(0.9 V^{t-1} + x @ w - th)` is read from `V^{t-1}` directly. The + perturbation at `V^t` cannot see that read. + +Measured on `main` (D-RTRL, `etrace_grad`, |grad| summed): | model | online | BPTT | |---|---|---| -| issue LIF, output `sg(V_pre - th)` (ETP weight through a pre-reset value) | 0.000000 | 1.079097 | +| issue LIF, output `sg(V_pre - th)` | 0.000000 | 1.079097 | | plain readout inside the step, `tanh(V) @ w_out` (`w_out`) | 0.0 | 16.27 | -The second row is the same defect, and probably the more common one: any -readout weight applied inside the compiled step gets no gradient. - -## Fix: add the exact same-step bypass term - -For a parameter `w` the true gradient decomposes, per step, into paths that -cross a new hidden value (handled by the trace, unchanged) and paths that do -not, the *bypass*. The bypass contribution is a plain same-step reverse-mode -term with no temporal component, so it is exact for every algorithm, exact or -approximate, and is added the same way for all of them. - -1. **Cone detection (compile time, perturbation pass).** Forward reachability - from every parameter leaf invar through the perturbed step. Propagation stops - at hidden-state output slots and at `stop_gradient`. A parameter whose cone - reaches a model output, or a non-hidden, non-param state output, is a - *bypass parameter*. An INFO diagnostic `OUTPUT_BYPASSES_HIDDEN` names them. -2. **Cone duplication (compile time).** For each bypass parameter the perturbed - jaxpr gets one extra invar per leaf (a differentiable copy of the value). - Every equation in the copies' cone is duplicated, consuming the copies - (or earlier duplicates) for cone inputs and `stop_gradient` of the original - variable for every other input. The original equations are kept unchanged, - so values, the hidden perturbation and every other gradient are untouched. - At each merge slot (model outputs and non-hidden state outputs) the - duplicate `d` is attached as a pure tangent: - `out = select_n(is_finite(d), out, out + (d - stop_gradient(d)))`. - The forward value is bitwise `out`, and the `select_n` guard keeps an - overflowing output from turning into NaN. Hidden slots and the extra - Jacobian temporaries are never merged. -3. **Executor.** In single-step mode, `etrace_param_vals` holds the bypass - parameters (instead of `{}`) and is fed to the copies only. The real - parameter values still come from the states, as constants. The VJP - therefore returns exactly the bypass gradient in `dg_etrace_params`. -4. **Backward.** `_update_fn_bwd` accepts `dg_etrace_params` whose keys are the - compiled bypass parameters (anything else still raises) and passes them as - `dl_to_etws_at_t`. Every `_solve_weight_gradients` already adds this term - (`param_dim_vjp`, `io_dim_vjp`, `random_projection_vjp`). For parameters - with no ETP relation it becomes their only gradient. Bootstrapped DNI is - unaffected: the future cotangent enters at hidden exits, which the cone - never crosses. -5. **Unfixable remainder: WARNING.** A cone equation that cannot be - duplicated is cut, with a WARNING (`OUTPUT_BYPASS_DROPPED`) naming the - parameters. That covers a `while` (not reverse-differentiable) or an - equation with effects (duplicating it would repeat the effect). - -No change when there is no bypass: the perturbed jaxpr and the executor are -exactly as before. Multi-step mode differentiates the parameters directly -and already includes the bypass. A test checks that, and that path is not -changed. +With only the parameter term, the LIF gives 0.234932. The hidden term +supplies the rest. + +## Fix: add the exact same-step bypass terms + +### Compile time (`_compiler/output_bypass.py`) + +`add_output_bypass(eqns, outvars, sources=..., cut_vars=..., merge_positions=...)` +rewrites the perturbed step. The sources are keyed `('param', path)` (the +parameter leaf invars) and `('hidden', path)` (the hidden invars, i.e. +`h^{t-1}`). + +1. **Cones.** Forward reachability from every source. Propagation stops at + the hidden-state output slots (`cut_vars`, whose same-step credit the + trace already carries) and at `stop_gradient`. A source whose cone + reaches a merge slot has a bypass. The merge slots are the model outputs + and the non-hidden, non-param state outputs; integer and boolean slots + are skipped. An INFO diagnostic `OUTPUT_BYPASSES_HIDDEN` names the + parameters and the hidden states. +2. **Duplication.** Each bypass source gets one new invar per leaf, a + differentiable copy. Every equation of the copies' cone is duplicated. + A duplicate consumes the copies (or earlier duplicates) for cone inputs, + and `stop_gradient` of the original variable for every other input. The + original equations are untouched, so values, the hidden perturbation and + every other gradient stay exactly as before. +3. **Merge.** At each merge slot the duplicate `d` is attached as a pure + tangent: `select_n(is_finite(d), o, o + (d - stop_gradient(d)))`. The + forward value is `o`, and an overflowing output stays `inf` rather than + becoming `inf - inf = nan`. +4. **Unfixable remainder: WARNING.** A cone equation that cannot be + duplicated is cut: a `while` anywhere inside it (no reverse rule) or an + equation with effects. `OUTPUT_BYPASS_DROPPED` names the sources. + +`HiddenPerturbation` gains `bypass_param_paths` / `bypass_vars`, +`bypass_hidden_paths` / `bypass_hidden_vars`, and +`eval_jaxpr(inputs, perturb_data, bypass_data=(), bypass_hidden_data=())`. +The new invars follow the perturbation invars. With no bypass, the +perturbed jaxpr is exactly as before. + +### Executor (`_algorithm/vjp_graph_executor.py`) + +In single-step mode, `etrace_param_vals` holds the bypass parameters and +feeds the copies only (the real parameter values still come from the states +as constants). A new VJP primal `hidden_bypass_vals` holds `h^{t-1}` of the +bypass hidden states. The VJP therefore returns the parameter bypass +gradient in `dg_etrace_params` and the hidden bypass signal in the new +cotangent slot. + +### Backward (`_algorithm/vjp_base.py`) + +- The single-step guard now requires `dg_etrace_params` to carry exactly the + compiled bypass parameters (anything else still raises). They pass on as + `dl_to_etws_at_t`, which every `_solve_weight_gradients` already adds. +- `_update_fn_fwd` keeps `e^{t-1}` in the residuals when there is a hidden + bypass. `_update_fn_bwd` turns the hidden bypass cotangents into per-group + arrays (zeros for group members without a bypass). It passes them through + `_compute_learning_signal` (so `random_feedback` projects them too) and + contracts them with `e^{t-1}` through a new pure hook, + `_solve_previous_trace_gradients`. The hook runs *before* the main + contraction, which may advance the ParamDim trace filters. +- The hook is implemented for each engine: + - ParamDim: the previous *filtered* trace when filters are active; + - IO-dim: the warm-up bias correction is indexed by `running_index - 1`; + - random projection: the same contraction as the main term. +- `learning_signal='modulatory'` replaces `dL/dh` entirely, so the graph is + compiled without the rewrite (`compile_etrace_graph(include_output_bypass=False)`). +- Bootstrapped DNI: the second (future) pass ignores the bypass slots, like + the perturbation slots. The learning signal comes from pass 1 only. + +Multi-step mode differentiates the parameters and the hidden values directly +and already includes both terms. It builds no perturbation and is unchanged. + +## Consequences + +- A same-step read of a pure hidden copy (`h2^t = h1^{t-1}`, #184) is now + credited exactly, as the hidden bypass of `h1^{t-1}`: the same value with + the same trace. The `HIDDEN_COPY_SAME_STEP_READ` warning added in #190 is + removed. +- A loss that reads a descended scan's stacked ys instead of its carry + (`scan_descent.py` note) now gets the same gradient as a carry read. The + limitation pin is changed to an equality test. ## Edge cases - Output reads the new hidden value: no bypass, jaxpr unchanged. - Mixed output (bypass + hidden path in one output): both terms, no double counting (the cone is cut at hidden slots). -- Weight feeds another ETP primitive on the readout path (`W2(W1 x)` with no - hidden state): the cone runs through `W2`'s primitive. The ETP JVP rules are - used, and both get bypass gradients. -- Output is the parameter itself (`return w`). -- Pytree / multi-leaf parameter (weight + bias). +- Two hidden states in one group, only one with a bypass: zero-filled. +- Weight feeds a second ETP op on the readout path (`W2(tanh(W1 x))`): the + cone runs through both primitives, and both get bypass gradients. +- A parameter applied directly to the output (`... + b`). +- A pytree parameter (`Linear` weight + bias). - Integer or boolean outputs: not merged. - Output overflow (`inf`): the forward value stays `inf`, not NaN. - `while` on the bypass path: dropped with a WARNING. -- Other (non-hidden) state written from a weight (`ShortTermState`): merged - slot, its incoming cotangent is honoured. -- vmap / batched learners. +- A non-hidden state written from a weight (`ShortTermState`): a merge slot, + so its same-step credit is exact. Recursion *through* such a state carries + no trace (it is not a hidden state), as before. +- vmap / per-lane learners. ## Tests -- `_compiler/hidden_pertubation_test.py`: - - detection: bypass and no-bypass; - - structure: no extra invars without a bypass; - - forward values bitwise unchanged; - - the gradient w.r.t. the copies equals a hand-written bypass VJP; - - `while` on the cone gives a WARNING; - - `inf` output stays `inf`. -- `_algorithm/d_rtrl_test.py`, `_algorithm/e_prop_test.py`, `pp_prop_test.py`: - - the issue's LIF and the plain-readout model: D-RTRL equals BPTT, - through `etrace_grad` (issue repro path) and through the finite-window - oracle with chunk size 1; - - mixed output: no double count; - - EProp / pp_prop: the bypass term is present and the algorithms keep - their own trace terms. Checked against the same algorithm on the - `new hidden value` twin model, via the finite-window oracle; - - multi-step mode already equals BPTT on these models (regression pin); - - batched (`batch_size=2`) learner. -- `_algorithm/vjp_base_test.py`: the guard still raises for unexpected ETP - gradients under single-step. +- `_compiler/output_bypass_test.py`: + - detection: the pre-reset output has a parameter and a hidden bypass; a + readout from the new hidden value has none, and its jaxpr is unchanged; + a plain readout has a parameter bypass only; `include_output_bypass=False`; + - `while` on the cone gives one WARNING naming both sources; + - forward values unchanged; gradients w.r.t. the copies equal the + hand-written partials; `inf` stays `inf`; length checks; + - `add_output_bypass` on hand-written jaxprs: `stop_gradient`, integer + outputs, cut vars, merge positions. +- `_algorithm/tests/output_bypass_test.py`: + - D-RTRL and EProp (exact on these element-wise cells) equal BPTT through + `etrace_grad` for: + - the pre-reset LIF and the mixed output; + - a two-state group; + - the pure-copy same-step read; + - the plain readout and the readout variants (ETP `Linear`, direct + parameter, ETP chain, non-hidden state); + - the issue's numbers; + - multi-step D-RTRL equals BPTT through the finite-window oracle (chunk 1 + and 4); + - ThreeFactor compiles without a bypass; + - pp_prop and UORO get both bypasses; their hooks equal the main + contraction of the previous trace; the ParamDim hook reads the filtered + trace without advancing it; + - vmapped lanes equal per-sample BPTT. +- `_algorithm/tests/scan_descent_support_test.py`: the ys readout equals the + carry readout. From a94217288206071d84b735fa6982a46b34954871 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 24 Oct 2026 18:27:09 +0800 Subject: [PATCH 3/5] Document outputs that bypass new hidden states (#185) --- docs/advanced/limitations.ipynb | 92 ++++++++++++++++++- .../2026-10-24-issue-184-input-only-hidden.md | 4 + 2 files changed, 93 insertions(+), 3 deletions(-) diff --git a/docs/advanced/limitations.ipynb b/docs/advanced/limitations.ipynb index 1fa3f71b..d22820df 100644 --- a/docs/advanced/limitations.ipynb +++ b/docs/advanced/limitations.ipynb @@ -395,6 +395,90 @@ } ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Outputs That Bypass New Hidden States\n", + "\n", + "Under the default `vjp_method='single-step'`, a parameter's online gradient is assembled from the eligibility trace, `dL/dh^t · e^t`. A step output can also be computed from a value that is **not** a new hidden state: a spike read from the pre-reset membrane potential, a readout weight applied inside the step, or the stacked ys of a descended scan. The compiler finds such *bypass* paths and adds their exact same-step credit:\n", + "\n", + "- the parameters' own gradient along the bypass (`bypass_param_paths`), and\n", + "- the learning signal `dL^t/dh^{t-1}` of a previous hidden value read along the bypass, paired with the previous trace `e^{t-1}` (`bypass_hidden_paths`).\n", + "\n", + "An INFO diagnostic `OUTPUT_BYPASSES_HIDDEN` names them. Below, the LIF output is read before the reset, and D-RTRL matches BPTT.\n", + "\n", + "A bypass path that crosses a `while_loop` or an equation with effects cannot be differentiated in reverse mode. That part is dropped with a WARNING (`OUTPUT_BYPASS_DROPPED`). `learning_signal='modulatory'` replaces `dL/dh` entirely and does not add the bypass terms." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "bypass parameters : [('w',)]\n", + "bypass hidden states: [('V',)]\n", + "D_RTRL |grad w| = 1.079097\n", + "BPTT |grad w| = 1.079097\n" + ] + } + ], + "source": [ + "import jax\n", + "import jax.numpy as jnp\n", + "import brainstate\n", + "import braintools\n", + "import braintrace\n", + "\n", + "sg = braintools.surrogate.Sigmoid()\n", + "\n", + "\n", + "class LIF(brainstate.nn.Module):\n", + " def __init__(self):\n", + " super().__init__()\n", + " self.w = brainstate.ParamState(brainstate.random.randn(3, 4))\n", + "\n", + " def init_state(self, *args, **kwargs):\n", + " self.V = brainstate.HiddenState(jnp.zeros(4))\n", + "\n", + " def update(self, x):\n", + " V_pre = 0.9 * self.V.value + braintrace.matmul(x, self.w.value)\n", + " spk = sg(V_pre - 1.0)\n", + " self.V.value = V_pre - jax.lax.stop_gradient(spk) # reset\n", + " return spk # read before the reset\n", + "\n", + "\n", + "def loss_fn(y, t):\n", + " return jnp.mean((y - t) ** 2)\n", + "\n", + "\n", + "brainstate.random.seed(42)\n", + "xs = brainstate.random.rand(10, 3) * 2.0\n", + "ys = jnp.full((10, 4), 0.3)\n", + "\n", + "brainstate.random.seed(0)\n", + "net = LIF()\n", + "brainstate.nn.init_all_states(net)\n", + "learner = braintrace.compile(net, braintrace.D_RTRL, xs[0])\n", + "print('bypass parameters :', list(learner.graph.hidden_perturb.bypass_param_paths))\n", + "print('bypass hidden states:', list(learner.graph.hidden_perturb.bypass_hidden_paths))\n", + "g_online = learner.etrace_grad(xs, ys, step_fn=lambda x, y: loss_fn(learner(x), y))\n", + "\n", + "brainstate.random.seed(0)\n", + "net = LIF()\n", + "brainstate.nn.init_all_states(net)\n", + "g_bptt = brainstate.transform.grad(\n", + " lambda: brainstate.transform.for_loop(lambda x, y: loss_fn(net(x), y), xs, ys).mean(),\n", + " grad_states={'w': net.w},\n", + ")()\n", + "print(f\"D_RTRL |grad w| = {float(jnp.abs(g_online[('w',)]).sum()):.6f}\")\n", + "print(f\"BPTT |grad w| = {float(jnp.abs(g_bptt['w']).sum()):.6f}\")" + ] + }, { "cell_type": "markdown", "id": "a1b2c3d4e5f60013", @@ -518,11 +602,13 @@ "\n", "5. **Do not write a `ParamState` inside the step.** Update trained weights outside the step from the `etrace_grad` gradients; hold locally-updated (plastic) weights in a non-param state such as `brainstate.LongTermState`.\n", "\n", - "6. **Choose the right algorithm based on memory/accuracy tradeoffs.** Use `D_RTRL` for exact gradients with moderate model sizes, and `ES_D_RTRL` for memory-efficient approximate gradients with larger models.\n", + "6. **Outputs may bypass new hidden states.** Under the single-step VJP, their exact same-step credit is added automatically (`OUTPUT_BYPASSES_HIDDEN`). The exception is a bypass that crosses a `while_loop`, which is dropped with a warning.\n", + "\n", + "7. **Choose the right algorithm based on memory/accuracy tradeoffs.** Use `D_RTRL` for exact gradients with moderate model sizes, and `ES_D_RTRL` for memory-efficient approximate gradients with larger models.\n", "\n", - "7. **Call `braintrace.compile` once before training**, not inside the training loop. The compiled graph is cached and reused for inputs of the same shape.\n", + "8. **Call `braintrace.compile` once before training**, not inside the training loop. The compiled graph is cached and reused for inputs of the same shape.\n", "\n", - "8. **The compiler works with all standard JAX mathematical operations.** As long as you avoid the unsupported control flow primitives listed above, your model will compile successfully." + "9. **The compiler works with all standard JAX mathematical operations.** As long as you avoid the unsupported control flow primitives listed above, your model will compile successfully." ] } ], diff --git a/docs/specs/2026-10-24-issue-184-input-only-hidden.md b/docs/specs/2026-10-24-issue-184-input-only-hidden.md index 6c402449..c9e57e13 100644 --- a/docs/specs/2026-10-24-issue-184-input-only-hidden.md +++ b/docs/specs/2026-10-24-issue-184-input-only-hidden.md @@ -67,6 +67,10 @@ is fixed here too. `HIDDEN_COPY_SAME_STEP_READ` that names both states and suggests reading the copy on the next step or computing the output from a new hidden value. + *Superseded by #185:* the output-bypass rewrite credits such a read as the + hidden bypass of `h1^{t-1}` (the same value, with the same trace + `e2^t = e1^{t-1}`), so the read is exact and the warning was removed. + ## Edge cases - `latch <- x` (no weight): the outvar is a top-level jaxpr invar. From 0fb7713e868548f1dc5f3d60411f8365965f897c Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 24 Oct 2026 18:35:35 +0800 Subject: [PATCH 4/5] Narrow F-33 to plain parameters that feed hidden state; type fixes (#185) --- braintrace/_algorithm/vjp_base.py | 11 +++++++-- braintrace/_compiler/hidden_pertubation.py | 17 +++++++------- braintrace/_compiler/output_bypass.py | 27 +++++++++++++--------- braintrace/_testing/oracle_models_test.py | 15 +++++++----- docs/specs/2026-07-25-known-limitations.md | 16 ++++++++----- 5 files changed, 53 insertions(+), 33 deletions(-) diff --git a/braintrace/_algorithm/vjp_base.py b/braintrace/_algorithm/vjp_base.py index 9cfc305c..3400b458 100644 --- a/braintrace/_algorithm/vjp_base.py +++ b/braintrace/_algorithm/vjp_base.py @@ -1313,6 +1313,12 @@ def _update_fn_bwd( dg_weights_prev = None if prev_etrace_vals is not None: hidden_perturb = self.graph.hidden_perturb + if hidden_perturb is None: + raise ValueError( + 'The forward pass kept the previous eligibility trace for a ' + 'hidden output bypass, but the compiled graph carries no ' + 'hidden perturbation. Recompile the graph.' + ) dl2h_prev = hidden_perturb.bypass_hidden_data_to_hidden_group_data( [u.get_mantissa(dg_hidden_bypass[path]) for path in hidden_perturb.bypass_hidden_paths], self.graph.hidden_groups, @@ -1843,11 +1849,12 @@ def _solve_previous_trace_gradients( dict Parameter path to its gradient (``None`` for no contribution). """ + hidden_perturb = self.graph.hidden_perturb + paths = [] if hidden_perturb is None else list(hidden_perturb.bypass_hidden_paths) raise NotImplementedError( f'{type(self).__name__} does not implement ' f'`_solve_previous_trace_gradients`, which the hidden output bypass ' - f'of the compiled graph (hidden states ' - f'{list(self.graph.hidden_perturb.bypass_hidden_paths)}) needs.' + f'of the compiled graph (hidden states {paths}) needs.' ) def _solve_weight_gradients( diff --git a/braintrace/_compiler/hidden_pertubation.py b/braintrace/_compiler/hidden_pertubation.py index 1f6f3644..52c5ade2 100644 --- a/braintrace/_compiler/hidden_pertubation.py +++ b/braintrace/_compiler/hidden_pertubation.py @@ -14,7 +14,7 @@ # ============================================================================== -from typing import Dict, FrozenSet, List, Optional, Set, Sequence, NamedTuple, Any +from typing import Dict, FrozenSet, List, Mapping, Optional, Set, Sequence, NamedTuple, Any, cast import brainstate import jax.core @@ -223,10 +223,10 @@ def bypass_hidden_data_to_hidden_group_data( f'got {len(bypass_hidden_data)}.' ) by_path = dict(zip(self.bypass_hidden_paths, bypass_hidden_data)) - per_perturb = [ - by_path[path] if path in by_path else jnp.zeros(v.aval.shape, v.aval.dtype) - for path, v in zip(self.perturb_hidden_paths, self.perturb_vars) - ] + per_perturb = [] + for path, v in zip(self.perturb_hidden_paths, self.perturb_vars): + aval = cast(jax.core.ShapedArray, v.aval) + per_perturb.append(by_path[path] if path in by_path else jnp.zeros(aval.shape, aval.dtype)) return self.perturb_data_to_hidden_group_data(per_perturb, hidden_groups) def perturb_data_to_hidden_group_data( @@ -372,7 +372,7 @@ def __init__( path_to_state: Dict[Path, brainstate.HiddenState], control_flow: ControlFlowPolicy = DEFAULT_CONTROL_FLOW_POLICY, descended_scan_eqn_ids: FrozenSet[int] = frozenset(), - weight_path_to_invars: Optional[Dict[Path, Sequence[Var]]] = None, + weight_path_to_invars: Optional[Mapping[Path, Sequence[Var]]] = None, bypass_merge_positions: Sequence[int] = (), ): # necessary data structures @@ -486,7 +486,8 @@ def compile(self) -> HiddenPerturbation: ) param_bypass_keys = [k for k in bypass.bypass_keys if k[0] == 'param'] hidden_bypass_keys = [k for k in bypass.bypass_keys if k[0] == 'hidden'] - n_param_leaves = sum(len(self.weight_path_to_invars[k[1]]) for k in param_bypass_keys) + # one copy per hidden bypass (a hidden invar is a single leaf) + n_param_leaves = len(bypass.bypass_vars) - len(hidden_bypass_keys) jaxpr = Jaxpr( constvars=list(self.closed_jaxpr.jaxpr.constvars), invars=( @@ -723,7 +724,7 @@ def add_hidden_perturbation_in_jaxpr( path_to_state: Dict[Path, brainstate.HiddenState], control_flow: ControlFlowPolicy = DEFAULT_CONTROL_FLOW_POLICY, descended_scan_eqn_ids: FrozenSet[int] = frozenset(), - weight_path_to_invars: Optional[Dict[Path, Sequence[Var]]] = None, + weight_path_to_invars: Optional[Mapping[Path, Sequence[Var]]] = None, bypass_merge_positions: Sequence[int] = (), ) -> HiddenPerturbation: """ diff --git a/braintrace/_compiler/output_bypass.py b/braintrace/_compiler/output_bypass.py index b42636c3..32030757 100644 --- a/braintrace/_compiler/output_bypass.py +++ b/braintrace/_compiler/output_bypass.py @@ -56,7 +56,7 @@ the effect). """ -from typing import Dict, FrozenSet, Hashable, List, NamedTuple, Sequence, Set, Tuple +from typing import Any, Dict, FrozenSet, Hashable, List, NamedTuple, Sequence, Set, Tuple, cast import jax import jax.numpy as jnp @@ -102,8 +102,9 @@ class OutputBypass(NamedTuple): eqns: List[JaxprEqn] outvars: List bypass_vars: List[Var] - bypass_keys: List[Hashable] - dropped_keys: List[Hashable] + # the caller's source keys (``Any``: callers index into structured keys) + bypass_keys: List[Any] + dropped_keys: List[Any] def _is_blocking(eqn: JaxprEqn) -> bool: @@ -115,8 +116,14 @@ def _is_blocking(eqn: JaxprEqn) -> bool: Jaxpr(constvars=[], invars=[], outvars=[], eqns=[eqn])) +def _shaped(v: Var) -> jax.core.ShapedArray: + # every variable of a traced step is array-valued + aval = cast(jax.core.ShapedArray, v.aval) + return jax.core.ShapedArray(aval.shape, aval.dtype) + + def _is_inexact(v: Var) -> bool: - return jnp.issubdtype(v.aval.dtype, jnp.inexact) + return jnp.issubdtype(_shaped(v).dtype, jnp.inexact) def _find_cones( @@ -211,7 +218,7 @@ def add_output_bypass( bypass_vars: List[Var] = [] for key in bypass_keys: for v in sources[key]: - copy = new_var('', jax.core.ShapedArray(v.aval.shape, v.aval.dtype)) + copy = new_var('', _shaped(v)) bypass_vars.append(copy) if v not in cut_vars: dup[v] = copy @@ -221,7 +228,7 @@ def add_output_bypass( def _detached(v: Var) -> Var: if v not in detached: - fresh = new_var('', jax.core.ShapedArray(v.aval.shape, v.aval.dtype)) + fresh = new_var('', _shaped(v)) new_eqns.append(new_jaxpr_eqn([v], [fresh], stop_gradient_p, {}, set())) detached[v] = fresh return detached[v] @@ -238,10 +245,7 @@ def _detached(v: Var) -> Var: v if isinstance(v, Literal) else (dup[v] if v in dup else _detached(v)) for v in eqn.invars ] - dup_outvars = [ - new_var('', jax.core.ShapedArray(ov.aval.shape, ov.aval.dtype)) - for ov in eqn.outvars - ] + dup_outvars = [new_var('', _shaped(ov)) for ov in eqn.outvars] new_eqns.append(eqn.replace(invars=dup_invars, outvars=dup_outvars)) for ov, dv in zip(eqn.outvars, dup_outvars): if ov not in cut_vars: @@ -251,7 +255,8 @@ def _detached(v: Var) -> Var: merged: Dict[Var, Var] = {} def _like(v: Var, dtype=None) -> Var: - return new_var('', jax.core.ShapedArray(v.aval.shape, dtype or v.aval.dtype)) + aval = _shaped(v) + return new_var('', jax.core.ShapedArray(aval.shape, dtype or aval.dtype)) for k in merge_positions: v = outvars[k] diff --git a/braintrace/_testing/oracle_models_test.py b/braintrace/_testing/oracle_models_test.py index 2d3fbc3a..fc62652a 100644 --- a/braintrace/_testing/oracle_models_test.py +++ b/braintrace/_testing/oracle_models_test.py @@ -283,21 +283,24 @@ def test_it_declares_both_plain_kinds(self): assert spec.etp_param_keys == (('w',),) assert spec.plain_param_keys == (('win',), ('wout',)) - def test_single_step_zeroes_every_plain_key(self): - # F-33. Not a truncation -- an absence. Pinned on the fixture that makes - # it visible, so a single-step rule's silence about plain parameters is a - # recorded property rather than a surprise. + def test_single_step_zeroes_plain_keys_that_feed_hidden_state(self): + # F-33. Not a truncation -- an absence. A plain parameter whose only + # path to the loss crosses a hidden state (``win``) has no trace, so a + # single-step rule is silent about it. A plain readout (``wout``) + # bypasses every new hidden value and gets its exact same-step + # gradient (#185), which is all of its credit. from braintrace._testing.oracle_models import plain_and_etp_rnn spec = plain_and_etp_rnn(n_in=3, n_rec=4, n_out=2) inputs = _inputs(4, 3) got = online_param_gradients_singlestep_naive( spec.factory, inputs, algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='single-step')) + ref = bptt_param_gradients(spec.factory, inputs) assert float(jnp.abs(got[('win',)]).max()) == 0.0 - assert float(jnp.abs(got[('wout',)]).max()) == 0.0 + np.testing.assert_allclose( + np.asarray(got[('wout',)]), np.asarray(ref[('wout',)]), atol=1e-5) assert float(jnp.abs(got[('w',)]).max()) > 1e-6 - class TestDelayedRewardRnn: """``delayed_reward_rnn`` -- credit that actually spans windows.""" diff --git a/docs/specs/2026-07-25-known-limitations.md b/docs/specs/2026-07-25-known-limitations.md index a7757942..c314eb52 100644 --- a/docs/specs/2026-07-25-known-limitations.md +++ b/docs/specs/2026-07-25-known-limitations.md @@ -58,7 +58,7 @@ resolution, never that a test stopped being run. | F-30 | the IO-dim f-side de-biasing correction is indexed by `update()` call count, not by trace-step count, so it is exact only for single-step input | active, **preserved deliberately** — fixing it moves `pp_prop` / `OSTLFeedforward` gradients | `io_dim_vjp_test.py::TestBiasCorrectionTimeIndex`, both the structural claim and the finite-window consequence | | F-31 | on a hidden group whose `y -> hidden` tail is not position-preserving, *no* within-group scope reaches BPTT — saturated SnAp / full within-group RTRL included | active, **structural** — pre-P3, affects every `recurrence_scope`; the compiler detects the tail and warns | `_algorithm/snap_n_test.py::TestConservativeFallbackEndToEnd::test_a_relabelling_tail_defeats_saturation_and_says_so`, against the `roll=0` control; measured below | | F-32 | the coupled-scope transition is materialised as a dense `(*V, S, *V, S)` Jacobian, so `random_projection` costs `O((V·S)^2)` per group even though it only ever needs `D @ s_tilde` — one JVP would do | active, **deliberate**; a matrix-free `D_full` would need a JVP seam the executor does not have | `_algorithm/vjp_graph_executor_test.py::TestFullJacobianFlag::test_the_flag_changes_the_jacobian_shape` records the shape; UORO's docstring states the memory honestly | -| F-33 | under `vjp_method='single-step'`, every **plain** (non-ETP) parameter's gradient is exactly **zero** — not merely truncated | active, **structural** | `oracle_models_test.py::test_single_step_zeroes_every_plain_key`, and for the preset that cannot avoid it, `three_factor_test.py::TestItDoesSomethingMeasurably::test_every_plain_parameter_gets_an_exact_zero_which_is_f_33`; consequence documented in `ThreeFactor`'s docstring | +| F-33 | under `vjp_method='single-step'`, a **plain** (non-ETP) parameter whose path to the loss crosses a hidden state gets exactly **zero** — not merely truncated (a plain readout is exact since #185) | active, **structural** | `oracle_models_test.py::test_single_step_zeroes_plain_keys_that_feed_hidden_state`, and for the preset that cannot avoid it, `three_factor_test.py::TestItDoesSomethingMeasurably::test_every_plain_parameter_gets_an_exact_zero_which_is_f_33`; consequence documented in `ThreeFactor`'s docstring | | F-34 | in the multi-step backward pass, `dg_etrace_params` carries the within-window gradient of **every** trainable parameter, plain ones included, while `dg_non_etrace_params` is empty — so the ETP/plain split cannot be read off which dictionary a parameter arrives in | active, **naming only**; the graph is the authority | `_algorithm/vjp_base_test.py::TestTheDefaultHooksAreInert::test_etp_routed_paths_reads_the_compiled_graph`; `_etp_routed_paths()` exists because of this | | F-35 | a `train_synthetic_gradient` fit is valid only for the exact `(loss_fn, chunk_size)` pair it was given, and **nothing checks that deployment matches** — a mismatch is not degraded DNI but noise shaped like a cotangent, measurably worse than leaving DNI off | active, **documented not enforced**; the learner never sees the caller's objective, so it has nothing to compare against. Narrowed since: the helper no longer *introduces* a mismatch of its own — a sequence length that is not a multiple of `chunk_size` used to be truncated into a short final window, fitting one pair against a shorter future than any deployed window, and is now refused | `dni_test.py::TestALearnedSynthesiserHelps::test_training_on_the_wrong_window_size_is_worse_than_not_training` pins the `chunk_size` half; the `loss_fn` half is measured in B4's docstring (0.577 against 0.140 for `M ≡ 0`). Since the sequence driver, `chunk_size` is also a keyword on `etrace_grad`, so fit-vs-deploy mismatch is reachable without writing a window by hand — `sequence_test.py::TestAlgorithmInteractions::test_a_chunk_size_mismatch_degrades_a_fitted_synthesiser` fits at 1, deploys at 2 (2.0e-02 relative), with an untrained-synthesiser control at 1.3e-01 so the comparison is not between two no-ops. `reduction` is **not** a mismatch surface: it divides once after the scan and never enters the differentiated objective (`|mean − sum/T| = 7.5e-09`) | | F-36 | DNI's pass-2 routing is **path**-granular, so a parameter used *both* as an ETP weight and plainly loses the future credit of its plain occurrence — the whole leaf is skipped because the path is ETP-routed | active, **structural**; the gradient dictionaries are keyed by path, not by occurrence, so occurrence-level routing is not expressible without changing every engine's `_solve_weight_gradients` contract | `dni_test.py::TestMixedRoutingIsPathGranular` | @@ -208,16 +208,20 @@ single-group model with an **elementwise** `y -> hidden` tail. ## Notes on F-33 -Under `vjp_method='single-step'` the residual jaxpr yields an **empty** -`dg_etrace_params` and an empty `dg_non_etrace_params`, so `dG_weights[path]` -is never written for a non-ETP parameter and stays `None` — which surfaces as an +Under `vjp_method='single-step'` the residual jaxpr yields only the +same-step *output-bypass* gradients in `dg_etrace_params` (#185): parameters +whose path to an output crosses no new hidden value. A non-ETP parameter whose +path to the loss crosses a hidden state has no trace and no bypass, so +`dG_weights[path]` is never written and stays `None` — which surfaces as an exact zero, not as a truncated approximation. Measured on `plain_and_etp_rnn` (ETP `w`, plain `win`, `wout`): the single-step total gradient is exactly zero -for both plain keys while `w` is finite. +for `win`, while the readout `wout` equals BPTT (all of its credit is +same-step) and `w` is finite. (Before #185 `wout` was zero too.) This is a property of the single-step window, not of any learning rule, but it bounds what a single-step-only rule can train. `ThreeFactor` is single-step by -construction (see its docstring for why), so a model given to it will train only +construction (see its docstring for why), and its modulatory signal replaces +`dL/dh`, so it compiles without the bypass: a model given to it will train only its ETP-routed parameters. The rule for callers: route what you intend to modulate through an ETP primitive. From b30ac1c6860b5e67dc10c7662da5e2b32cde4852 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 24 Oct 2026 18:51:34 +0800 Subject: [PATCH 5/5] Relative tolerance for the multi-step readout oracle test --- braintrace/_algorithm/tests/output_bypass_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/braintrace/_algorithm/tests/output_bypass_test.py b/braintrace/_algorithm/tests/output_bypass_test.py index ad3de57c..1bba6b96 100644 --- a/braintrace/_algorithm/tests/output_bypass_test.py +++ b/braintrace/_algorithm/tests/output_bypass_test.py @@ -249,7 +249,8 @@ def factory(): online = chunked_online_param_gradients( factory, xs, chunk_size=chunk_size, algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) - assert_param_gradients_close(online, bptt_param_gradients(factory, xs), atol=1e-5) + # gradients are O(10): float32 summation order differs across JAX versions + assert_param_gradients_close(online, bptt_param_gradients(factory, xs), atol=1e-5, rtol=1e-5) def test_modulatory_signal_compiles_without_bypass():