braintrace 0.3.0__tar.gz → 0.3.1__tar.gz

{braintrace-0.3.0 → braintrace-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.3.0
+Version: 0.3.1
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chaobrain@outlook.com>
 License: Apache-2.0 license

{braintrace-0.3.0 → braintrace-0.3.1}/braintrace/_algorithm/e_prop_test.py

@@ -15,9 +15,13 @@
 
 import unittest
 
+import brainevent
 import brainstate
 import jax
 import jax.numpy as jnp
+import numpy as np
+import numpy.testing as npt
+import pytest
 
 import braintrace
 from braintrace._algorithm.e_prop import EProp
@@ -443,3 +447,89 @@ def test_docstring_compile_example_runs():
     assert y.shape == (1,)
     assert bool(jnp.all(jnp.isfinite(y)))
     assert len(learner.graph.hidden_param_op_relations) >= 1
+
+
+# ----------------------------------------------------------------------------
+# #257: the kappa trace filter under compile(..., vmap=True)
+# ----------------------------------------------------------------------------
+
+_B, _N, _T = 3, 5, 6
+_rng = np.random.RandomState(0)
+_DENSE = ((_rng.rand(_N, _N) < 0.5) * _rng.randn(_N, _N)).astype('float32')
+_XS = jnp.asarray(_rng.randn(_T, _B, _N).astype('float32'))
+
+
+class _FilterNet(brainstate.nn.Module):
+    """The issue's model: a leaky recurrent layer through one ETP weight."""
+
+    def __init__(self, kind):
+        super().__init__()
+        if kind == 'sparse':
+            self.w = braintrace.nn.SparseLinear(
+                brainevent.CSR.fromdense(jnp.asarray(_DENSE)))
+        else:
+            self.w = braintrace.nn.Linear(_N, _N, w_init=jnp.asarray(_DENSE))
+
+    def init_state(self, *args, **kwargs):
+        self.h = brainstate.HiddenState(jnp.zeros(_N))
+
+    def update(self, x):
+        self.h.value = 0.9 * self.h.value + self.w(jnp.tanh(self.h.value) + x)
+        return self.h.value
+
+
+def _filtered_grad(kind, kappa, xs, **kw):
+    learner = braintrace.compile(_FilterNet(kind), 'eprop', xs[0],
+                                 kappa_filter_decay=kappa, **kw)
+    g = learner.etrace_grad(xs, step_fn=lambda x: jnp.sum(learner(x) ** 2))
+    return learner, g
+
+
+@pytest.mark.parametrize('vjp_method', ['single-step', 'multi-step'])
+@pytest.mark.parametrize('kind', ['dense', 'sparse'])
+def test_kappa_filter_under_vmap_is_the_sum_of_unbatched_lanes(kind, vjp_method):
+    # The batched objective sums the lanes' losses, so its gradient of the
+    # shared weights is the sum of the B unbatched gradients (f32 tolerance:
+    # the batched path reassociates the lane sum).
+    batched, g_b = _filtered_grad(kind, 0.5, _XS, batch_size=_B, vmap=True,
+                                  vjp_method=vjp_method)
+    lanes = [_filtered_grad(kind, 0.5, _XS[:, b], vjp_method=vjp_method)
+             for b in range(_B)]
+    g_ref = jax.tree.map(lambda *g: sum(g), *[g for _, g in lanes])
+    jax.tree.map(lambda a, r: npt.assert_allclose(a, r, rtol=1e-5, atol=1e-5),
+                 g_b, g_ref)
+    # the filter contributes: kappa=0 gives a different gradient
+    _, g_0 = _filtered_grad(kind, 0.0, _XS, batch_size=_B, vmap=True,
+                            vjp_method=vjp_method)
+    diff = max(float(jnp.max(jnp.abs(a - b)))
+               for a, b in zip(jax.tree.leaves(g_b), jax.tree.leaves(g_0)))
+    assert diff > 1e-2
+    # each lane's filter state is that lane's unbatched filter
+    flt_b = [f.value for f in batched.module._trace_filters.values()]
+    for b, (lane, _) in enumerate(lanes):
+        flt_l = [f.value for f in lane._trace_filters.values()]
+        jax.tree.map(lambda a, r: npt.assert_allclose(a[b], r, rtol=1e-5, atol=1e-6),
+                     flt_b, flt_l)
+
+
+@pytest.mark.parametrize('vjp_method', ['single-step', 'multi-step'])
+def test_kappa_filter_advances_in_the_forward_pass(vjp_method):
+    # A step that is not differentiated advances the filter exactly as a
+    # differentiated one does: the filter is a function of the trace sequence.
+    learners = []
+    for differentiate in (False, True):
+        learner = braintrace.compile(_FilterNet('dense'), 'eprop', _XS[0, 0],
+                                     kappa_filter_decay=0.5, vjp_method=vjp_method)
+
+        def step(x, learner=learner):
+            return jnp.sum(learner(x) ** 2)
+
+        for t in range(3):
+            if differentiate:
+                brainstate.transform.grad(step, learner.param_states)(_XS[t, 0])
+            else:
+                step(_XS[t, 0])
+        learners.append(learner)
+    plain, diffed = ([f.value for f in l._trace_filters.values()] for l in learners)
+    jax.tree.map(npt.assert_array_equal, plain, diffed)
+    assert max(float(jnp.max(jnp.abs(v))) for v in jax.tree.leaves(plain)) > 1e-2

{braintrace-0.3.0 → braintrace-0.3.1}/braintrace/_algorithm/param_dim_vjp.py

@@ -1553,6 +1553,60 @@ class ParamDimVjpAlgorithm(ETraceVjpAlgorithm):
         for x, val in etrace_vals.items():
             self.etrace_bwg[x].value = val
 
+    def _get_trace_filter_data(self) -> Optional[Dict[ETraceWG_Key, PyTree]]:
+        """The ``trace_filter='kappa'`` values, or ``None`` without a filter."""
+        if not self._trace_filters:
+            return None
+        return {k: v.value for k, v in self._trace_filters.items()}
+
+    def _assign_trace_filter_data(self, filter_vals: Optional[Dict[ETraceWG_Key, PyTree]]) -> None:
+        """Store the advanced ``trace_filter='kappa'`` values."""
+        if filter_vals is None:
+            return
+        for key, val in filter_vals.items():
+            self._trace_filters[key].value = val
+
+    def _advance_trace_filter(
+        self,
+        filter_vals: Optional[Dict[ETraceWG_Key, PyTree]],
+        etrace_vals: Dict[ETraceWG_Key, PyTree],
+    ) -> Tuple[Dict[ETraceWG_Key, PyTree], Optional[Dict[ETraceWG_Key, PyTree]]]:
+        r"""Low-pass the raw trace before it is contracted with the signal.
+
+        Applied elementwise (``jax.tree.map`` over the trace's own PyTree,
+        including its trailing ``num_state`` axis), so multi-state HiddenGroups
+        are filtered independently per state -- never summed across states and
+        broadcast back:
+
+        .. math::
+
+            \bar e^t_{ji} = \kappa\,\bar e^{t-1}_{ji} + e^t_{ji}
+
+        Returns the filtered trace, for the backward pass, and the new filter.
+        """
+        if filter_vals is None:
+            return etrace_vals, None
+        kappa = self.config.kappa
+        filtered = dict(etrace_vals)
+        new_filter = dict(filter_vals)
+        for key, trace in etrace_vals.items():
+            if key not in filter_vals:
+                continue
+            new_val = jax.tree.map(lambda prev, e: kappa * prev + e, filter_vals[key], trace)
+            filtered[key] = new_val
+            new_filter[key] = new_val
+        return filtered, new_filter
+
+    def _previous_filtered_trace(
+        self,
+        filter_vals: Optional[Dict[ETraceWG_Key, PyTree]],
+        etrace_vals: Dict[ETraceWG_Key, PyTree],
+    ) -> Dict[ETraceWG_Key, PyTree]:
+        """The previous *filtered* trace (the filter before this step)."""
+        if filter_vals is None:
+            return etrace_vals
+        return {k: filter_vals.get(k, v) for k, v in etrace_vals.items()}
+
     def _make_scan_fn(self, weight_vals: Dict[Path, PyTree]) -> Callable:
         """Build the per-step D-RTRL eligibility-trace stepper (scan body).
 
@@ -1668,17 +1722,13 @@ class ParamDimVjpAlgorithm(ETraceVjpAlgorithm):
         """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.
+        trace filters active, the forward pass hands in the previous *filtered*
+        trace (the filter before this step; see
+        :meth:`_previous_filtered_trace`).
         """
-        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,
+            etrace_h2w_at_t_minus_1,
             dG_weights,
             dl_to_hidden_groups,
             self.graph.hidden_param_op_relations,
@@ -1713,7 +1763,9 @@ class ParamDimVjpAlgorithm(ETraceVjpAlgorithm):
             Current timestep counter used for bias correction.
         etrace_h2w_at_t : Dict[ETraceWG_Key, PyTree]
             Eligibility trace values at the current timestep, mapping parameter-hidden
-            state relationship keys to trace values.
+            state relationship keys to trace values. Under ``trace_filter='kappa'``
+            the forward pass has already filtered them
+            (:meth:`_advance_trace_filter`).
         dl_to_hidden_groups : Sequence[jax.Array]
             Gradients of the loss with respect to hidden states at the current timestep.
         weight_vals : Dict[Path, PyTree]
@@ -1729,29 +1781,6 @@ class ParamDimVjpAlgorithm(ETraceVjpAlgorithm):
         Dict[Path, PyTree]
             Dictionary mapping parameter paths to their gradient values.
         """
-        # `trace_filter`: low-pass the raw trace before it is contracted with the
-        # learning signal. Applied elementwise (`jax.tree.map` over the trace's
-        # own PyTree, including its trailing num_state axis) so multi-state
-        # HiddenGroups are filtered independently per state -- never summed
-        # across states and broadcast back.
-        #
-        #     e_bar^t_{ji} = kappa * e_bar^{t-1}_{ji} + e^t_{ji}
-        #
-        if self._trace_filters:
-            kappa = self.config.kappa
-            filtered: Dict[ETraceWG_Key, PyTree] = {}
-            for key, trace in etrace_h2w_at_t.items():
-                flt = self._trace_filters.get(key)
-                if flt is None:
-                    filtered[key] = trace
-                    continue
-                new_val = jax.tree.map(
-                    lambda prev, e: kappa * prev + e, flt.value, trace
-                )
-                flt.value = new_val
-                filtered[key] = new_val
-            etrace_h2w_at_t = filtered
-
         dG_weights: Dict[Path, Any] = {path: None for path in self.param_states}
 
         # update the etrace weight gradients

{braintrace-0.3.0 → braintrace-0.3.1}/braintrace/_algorithm/sequence.py

@@ -573,3 +573,37 @@ class ETraceVmap(SequenceDriverMixin, brainstate.nn.Vmap):
     @property
     def _seq_vjp_method(self) -> Optional[str]:
         return getattr(self.module, 'vjp_method', None)
+
+    def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None:
+        """Reset every lane's eligibility traces and the running index.
+
+        ``brainstate.nn.Vmap`` inherits the no-op ``Module.reset_state``, so
+        without this a batched learner kept the previous sequence's traces and
+        every call after the first started from them. The wrapped learner's
+        ``reset_state`` zeroes its traces with ``zeros_like``, which keeps the
+        lane axis that ``compile(..., vmap=True)`` gave them, so this delegates
+        to it. Like the learner's own ``reset_state``, it leaves the model's
+        hidden states alone.
+
+        Parameters
+        ----------
+        batch_size : int, optional
+            Must be ``None``: the lane axis is fixed by
+            ``compile(..., batch_size=...)``, and a batch size here would put a
+            second batch axis in front of it.
+        **kwargs
+            Passed to the wrapped learner's ``reset_state``.
+
+        Raises
+        ------
+        ValueError
+            If ``batch_size`` is not ``None``.
+        """
+        if batch_size is not None:
+            raise ValueError(
+                f'ETraceVmap.reset_state(batch_size={batch_size!r}): the lane '
+                f'axis is fixed by compile(..., vmap=True, batch_size=...), and '
+                f'the traces already carry it; call reset_state() without '
+                f'batch_size.'
+            )
+        self.module.reset_state(**kwargs)  # type: ignore[attr-defined]

{braintrace-0.3.0 → braintrace-0.3.1}/braintrace/_algorithm/sequence_test.py

@@ -775,6 +775,63 @@ class TestVmap:
         for k, v in _arrays(grads).items():
             assert np.all(np.isfinite(v)), f'{k} is not finite'
 
+    def test_reset_state_zeroes_every_lanes_traces(self):
+        """``ETraceVmap.reset_state`` used to be ``Module``'s no-op.
+
+        After a call the traces are non-zero; ``reset_state`` must zero them
+        all, keep their lane axis and rewind ``running_index``.
+        """
+        batch = 3
+        xs, ys = _lane_data(batch)
+        learner = _vmap_learner(batch)
+        learner.etrace_grad(
+            xs, ys, step_fn=lambda i, t: jnp.sum((learner(i) - t) ** 2))
+        before = {k: s.value for k, s in learner.module.etrace_bwg.items()}
+        assert max(float(jnp.max(jnp.abs(x))) for x in jax.tree.leaves(before)) > 0.0
+        assert int(learner.module.running_index.value) == T
+
+        learner.reset_state()
+
+        for k, s in learner.module.etrace_bwg.items():
+            for old, new in zip(jax.tree.leaves(before[k]), jax.tree.leaves(s.value)):
+                assert jnp.shape(new) == jnp.shape(old)
+                assert jnp.shape(new)[0] == batch
+                assert float(jnp.max(jnp.abs(new))) == 0.0, f'{k} kept its trace'
+        assert int(learner.module.running_index.value) == 0
+
+    def test_a_second_call_after_reset_state_repeats_the_first(self):
+        """Two consecutive batched calls, with ``reset_state`` and the hidden
+        states restored between them, give the same gradient. Before the fix
+        the second call started from the first call's traces.
+        """
+        batch = 3
+        xs, ys = _lane_data(batch)
+        net = _VmapNet()
+        learner = braintrace.compile(net, 'D_RTRL', jnp.zeros((batch, N_IN)),
+                                     batch_size=batch, vmap=True,
+                                     vjp_method='multi-step')
+        hidden0 = {k: s.value for k, s in net.states(brainstate.HiddenState).items()}
+
+        def step_fn(inp, tar):
+            return jnp.sum((learner(inp) - tar) ** 2)
+
+        g_first = learner.etrace_grad(xs, ys, step_fn=step_fn)
+        learner.reset_state()
+        for k, s in net.states(brainstate.HiddenState).items():
+            s.value = hidden0[k]
+        g_second = learner.etrace_grad(xs, ys, step_fn=step_fn)
+
+        for k, v in _arrays(g_first).items():
+            assert np.max(np.abs(v)) > 0.0, f'{k} is identically zero -- vacuous'
+        _assert_trees_equal(g_second, g_first, rtol=1e-6, atol=1e-7,
+                            msg='second call after reset_state vs the first')
+
+    def test_reset_state_refuses_a_batch_size(self):
+        """The lane axis is the vmap's; a ``batch_size`` would add a second one."""
+        learner = _vmap_learner(3)
+        with pytest.raises(ValueError, match='batch_size'):
+            learner.reset_state(3)
+
     def test_the_vmap_return_value_is_still_a_brainstate_vmap(self):
         """Spec test 20 -- existing ``vmap=True`` users must be unaffected."""
         learner = _vmap_learner(3)

{braintrace-0.3.0 → braintrace-0.3.1}/braintrace/_algorithm/tests/output_bypass_test.py

@@ -35,6 +35,7 @@ import brainstate
 import braintools
 import jax
 import jax.numpy as jnp
+import numpy as np
 import numpy.testing as npt
 import pytest
 
@@ -303,22 +304,69 @@ def test_previous_trace_hook_is_the_weight_contraction(algo):
 
 
 def test_previous_trace_hook_reads_the_filtered_trace_without_advancing_it():
+    # The forward pass hands the bypass the previous *filtered* trace -- the
+    # filter before this step -- and neither it nor the contraction advances
+    # the filter (it advances once per step, in the forward pass; #257).
     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)
+    prev_trace = learner._previous_filtered_trace(learner._get_trace_filter_data(), trace)
+    assert set(prev_trace) == set(trace)
+    for k in trace:
+        assert prev_trace[k] is before[k]
+    prev = learner._solve_previous_trace_gradients(0, prev_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)
+    # the filter differs from the raw trace, so the pairing matters
+    raw = learner._solve_previous_trace_gradients(0, trace, dl, weights)
+    assert any(not np.allclose(prev[p], raw[p]) for p in weights)
+    # and the forward pass is what pairs it: its bypass residual is the filter
+    # before the step, its main residual the filter advanced on the new trace
+    xs, _ = _data(N)
+    fwd_out, fwd_res = learner._update_fn_fwd(
+        learner._conform_inputs((xs[4],)),
+        weights,
+        {k: s.value for k, s in learner.hidden_states.items()},
+        {k: s.value for k, s in learner.other_states.items()},
+        (trace, learner._get_trace_filter_data()),
+        learner.running_index.value,
+        None,
+    )
+    new_trace, new_filter = fwd_out[-1]
+    kappa = learner.config.kappa
+    for k in trace:
+        jax.tree.map(npt.assert_array_equal, fwd_res[-1][k], before[k])
+        jax.tree.map(lambda f, b, e: npt.assert_allclose(f, kappa * b + e, rtol=1e-6),
+                     fwd_res[1][k], before[k], new_trace[k])
+        jax.tree.map(npt.assert_array_equal, new_filter[k], fwd_res[1][k])
+
+
+def test_kappa_filter_with_bypass_under_vmap_is_the_sum_of_unbatched_lanes():
+    # e-prop's trace filter on a model with a hidden output bypass, under
+    # compile(vmap=True) (#257): the bypass pairs the filter before the step,
+    # which the forward pass hands the backward pass. The batched objective
+    # sums the lanes' losses, so the gradient is the sum of the B unbatched
+    # lanes' (f32 tolerance: the batched path reassociates the lane sum).
+    B, kappa = 3, 0.7
+    xs = jax.random.uniform(jax.random.PRNGKey(5), (T, B, N_IN)) * 2.0
+
+    def grad(x_seq, **kw):
+        brainstate.random.seed(0)
+        learner = braintrace.compile(_LIF('pre_reset'), 'eprop', x_seq[0],
+                                     kappa_filter_decay=kappa, **kw)
+        module = learner.module if kw.get('vmap') else learner
+        assert list(module.graph.hidden_perturb.bypass_hidden_paths) == [('V',)]
+        return learner.etrace_grad(
+            x_seq, step_fn=lambda x: jnp.sum((learner(x) - 0.3) ** 2))
+
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore')
+        batched = grad(xs, batch_size=B, vmap=True)
+        lanes = [grad(xs[:, b]) for b in range(B)]
+    ref = jax.tree.map(lambda *g: sum(g), *lanes)
+    jax.tree.map(lambda a, r: npt.assert_allclose(a, r, rtol=1e-5, atol=1e-5), batched, ref)
+    assert max(float(jnp.max(jnp.abs(g))) for g in jax.tree.leaves(ref)) > 1e-2
 
 
 def test_vmapped_lanes_match_per_sample_bptt():