braintrace 0.2.6__tar.gz → 0.2.7__tar.gz

{braintrace-0.2.6 → braintrace-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.6
+Version: 0.2.7
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -13,20 +13,20 @@ Classifier: Natural Language :: English
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries
-Requires-Python: >=3.11
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: jax>=0.8.0
+Requires-Dist: jax>=0.9.0
 Requires-Dist: brainstate>=0.5.2
 Requires-Dist: brainunit
 Requires-Dist: brainevent>=0.1.2

{braintrace-0.2.6 → braintrace-0.2.7}/braintrace/_algorithm/base.py

@@ -21,7 +21,7 @@
 
 from __future__ import annotations
 
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Tuple, cast
 
 import brainstate
 
@@ -257,11 +257,16 @@ class ETraceAlgorithm(SequenceDriverMixin, brainstate.nn.Module):
         #   - the weight states, which is invariant during the training process
         #   - the hidden states, the recurrent states, which may be changed between different training epochs
         #   - the other states, which may be changed between different training epochs
+        # ``FlattedDict.split`` returns one ``FlattedDict`` per filter (the
+        # stubs only promise a ``NestedDict``)
         (
             self._param_states,
             self._hidden_states,
             self._other_states
-        ) = self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...)
+        ) = cast(
+            Tuple[brainstate.util.FlattedDict, brainstate.util.FlattedDict, brainstate.util.FlattedDict],
+            self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...),
+        )
 
     def compile_graph(self, *args: Any) -> None:
         r"""

{braintrace-0.2.6 → braintrace-0.2.7}/braintrace/_algorithm/d_rtrl_test.py

@@ -1112,3 +1112,93 @@ class TestCoupledTraceBoundedness:
         assert all(bool(jnp.all(jnp.isfinite(leaf))) for leaf in leaves)
         max_abs = max(float(jnp.max(jnp.abs(leaf))) for leaf in leaves)
         assert max_abs < 1e3, f'eligibility trace grew to {max_abs}'
+
+
+# ---------------------------------------------------------------------------
+# Hidden-to-hidden edges without an equation (#184)
+# ---------------------------------------------------------------------------
+
+class _LatchModel(brainstate.nn.Module):
+    """``latch`` depends only on the input; ``V`` reads the old latch.
+
+    ``workaround=True`` writes ``latch <- new + 0 * latch`` (the pre-fix
+    workaround), which keeps a hidden-to-hidden edge in the jaxpr.
+    """
+
+    def __init__(self, workaround=False, n_in=3, n=4):
+        super().__init__()
+        self.workaround = workaround
+        self.n = n
+        self.w = brainstate.ParamState(0.5 * jax.random.normal(jax.random.PRNGKey(0), (n_in, n)))
+
+    def init_state(self, batch_size=None, **kwargs):
+        shape = (self.n,) if batch_size is None else (batch_size, self.n)
+        self.latch = brainstate.HiddenState(jnp.zeros(shape))
+        self.V = brainstate.HiddenState(jnp.zeros(shape))
+
+    def update(self, x):
+        I = self.latch.value
+        new = braintrace.matmul(x, self.w.value)
+        self.latch.value = new + 0 * self.latch.value if self.workaround else new
+        self.V.value = 0.9 * self.V.value + I
+        return jnp.tanh(self.V.value)
+
+
+class _DelayLineModel(brainstate.nn.Module):
+    """Two-step delay line ``h2^t = h1^{t-1}`` (a pure copy) feeding ``V``."""
+
+    def __init__(self, workaround=False, n_in=3, n=4):
+        super().__init__()
+        self.workaround = workaround
+        self.n = n
+        self.w = brainstate.ParamState(0.5 * jax.random.normal(jax.random.PRNGKey(1), (n_in, n)))
+
+    def init_state(self, batch_size=None, **kwargs):
+        shape = (self.n,) if batch_size is None else (batch_size, self.n)
+        self.h1 = brainstate.HiddenState(jnp.zeros(shape))
+        self.h2 = brainstate.HiddenState(jnp.zeros(shape))
+        self.V = brainstate.HiddenState(jnp.zeros(shape))
+
+    def update(self, x):
+        old1, old2 = self.h1.value, self.h2.value
+        self.V.value = 0.9 * self.V.value + old2
+        self.h1.value = 0.5 * old1 + braintrace.matmul(x, self.w.value)
+        self.h2.value = old1 + 0 * old2 if self.workaround else old1
+        return jnp.tanh(self.V.value)
+
+
+class TestEquationFreeHiddenEdges:
+    """Input-only latches and pure copies train exactly under D-RTRL (#184)."""
+
+    xs = jax.random.uniform(jax.random.PRNGKey(7), (8, 1, 3))
+
+    @pytest.mark.parametrize('model_cls', [_LatchModel, _DelayLineModel])
+    @pytest.mark.parametrize('chunk_size', [1, 3])
+    def test_d_rtrl_matches_bptt(self, model_cls, chunk_size):
+        from braintrace._testing.oracle import (
+            assert_model_is_live, assert_param_gradients_close,
+            bptt_param_gradients, chunked_online_param_gradients,
+        )
+        assert_model_is_live(model_cls, self.xs)
+        online = chunked_online_param_gradients(
+            model_cls, self.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(model_cls, self.xs), atol=1e-5)
+
+    @pytest.mark.parametrize('model_cls', [_LatchModel, _DelayLineModel])
+    @pytest.mark.parametrize('algo', ['d_rtrl', 'eprop'])
+    def test_matches_workaround(self, model_cls, algo):
+        xs = self.xs[:, 0]
+        ys = jnp.full((xs.shape[0], 4), 0.3)
+
+        def grads(workaround):
+            net = model_cls(workaround)
+            brainstate.nn.init_all_states(net)
+            learner = braintrace.compile(net, algo, xs[0])
+            return learner.etrace_grad(
+                xs, ys, step_fn=lambda x, y: jnp.mean((learner(x) - y) ** 2))
+
+        fixed, ref = grads(False), grads(True)
+        for k in ref:
+            assert float(jnp.abs(ref[k]).sum()) > 1e-6
+            npt.assert_allclose(fixed[k], ref[k], atol=1e-6)

{braintrace-0.2.6 → braintrace-0.2.7}/braintrace/_algorithm/io_dim_vjp.py

@@ -34,6 +34,7 @@ import jax
 import jax.numpy as jnp
 import brainunit as u
 
+from braintrace._compatible_imports import shaped_aval
 from braintrace._compiler import ControlFlowPolicy, HiddenGroup, HiddenParamOpRelation
 from braintrace._op import (
     etp_elemwise_p,
@@ -262,8 +263,8 @@ def _init_IO_dim_state(
         if x_key not in etrace_xs:
             x_repr_fn = get_pp_x_repr(relation.primitive)
             if x_repr_fn is None:
-                shape = relation.x_var.aval.shape
-                dtype = relation.x_var.aval.dtype
+                shape = shaped_aval(relation.x_var).shape
+                dtype = shaped_aval(relation.x_var).dtype
             else:
                 # the trace filters the primitive's x *representation*
                 # (e.g. embedding: the one-hot encoding of its integer
@@ -272,12 +273,12 @@ def _init_IO_dim_state(
                 x_aval = jax.eval_shape(
                     lambda x_, _fn=x_repr_fn, _w=weight_avals: _fn(x_, _w),
                     jax.ShapeDtypeStruct(
-                        relation.x_var.aval.shape, relation.x_var.aval.dtype),
+                        shaped_aval(relation.x_var).shape, shaped_aval(relation.x_var).dtype),
                 )
                 shape, dtype = x_aval.shape, x_aval.dtype
             etrace_xs[x_key] = EligibilityTrace(u.math.zeros(shape, dtype))
 
-    y_shape = relation.y_var.aval.shape
+    y_shape = shaped_aval(relation.y_var).shape
     group: HiddenGroup
     for group in relation.hidden_groups:
         # Exact match required, or (elemwise only) allow trailing-dim match
@@ -613,7 +614,7 @@ def _solve_IO_dim_weight_gradients(
             # gate the branch.)
             elemwise_batched = (
                 relation.primitive is etp_elemwise_p
-                and len(group.varshape) > relation.y_var.aval.ndim
+                and len(group.varshape) > shaped_aval(relation.y_var).ndim
             )
 
             if fast_solve:
@@ -926,8 +927,9 @@ class IODimVjpAlgorithm(ETraceVjpAlgorithm):
             find_this_weight = True
 
             # get the weight_op input
-            wx_var = etrace_x_key(relation.x_var)
-            if wx_var is not None:
+            # (an element-wise relation has no x trace)
+            if relation.x_var is not None:
+                wx_var = etrace_x_key(relation.x_var)
                 etrace_xs[wx_var] = self.etrace_xs[wx_var].value
 
             # get the weight_op df
@@ -1106,6 +1108,35 @@ class IODimVjpAlgorithm(ETraceVjpAlgorithm):
 
         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,

{braintrace-0.2.6 → braintrace-0.2.7}/braintrace/_algorithm/io_dim_vjp_test.py

@@ -247,6 +247,27 @@ class TestStateLifecycle:
         assert isinstance(result, tuple)
         assert len(result) == 2
 
+    def test_get_etrace_of_elementwise_weight(self):
+        # An element-wise relation has no x trace; only its df traces are returned.
+        class Leaky(brainstate.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = brainstate.ParamState(jnp.full((3,), 0.5))
+
+            def init_state(self, *args, **kwargs):
+                self.h = brainstate.HiddenState(jnp.zeros(3))
+
+            def update(self, x):
+                self.h.value = 0.9 * self.h.value + braintrace.element_wise(self.w.value) * x
+                return self.h.value
+
+        model = Leaky()
+        brainstate.nn.init_all_states(model)
+        algo = _compiled(model)
+        xs, dfs = algo.get_etrace_of(model.w)
+        assert xs == {}
+        assert len(dfs) == 1
+
     def test_get_etrace_of_plain_weight_raises(self):
         model = _build(om.tanh_rnn)
         algo = _compiled(model)

{braintrace-0.2.6 → braintrace-0.2.7}/braintrace/_algorithm/param_dim_vjp.py

@@ -1408,6 +1408,35 @@ class ParamDimVjpAlgorithm(ETraceVjpAlgorithm):
 
         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,

{braintrace-0.2.6 → braintrace-0.2.7}/braintrace/_algorithm/random_projection_vjp.py

@@ -665,6 +665,36 @@ class RandomProjectionVjpAlgorithm(ETraceVjpAlgorithm):
     # 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 @@ class RandomProjectionVjpAlgorithm(ETraceVjpAlgorithm):
         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)