braintrace 0.2.3__tar.gz → 0.2.5__tar.gz

{braintrace-0.2.3 → braintrace-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.3
+Version: 0.2.5
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -26,9 +26,10 @@ Classifier: Topic :: Software Development :: Libraries
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: jax>=0.8.0
 Requires-Dist: brainstate>=0.5.2
 Requires-Dist: brainunit
-Requires-Dist: brainevent
+Requires-Dist: brainevent>=0.1.2
 Requires-Dist: brainpy.state
 Requires-Dist: braintools
 Provides-Extra: cpu
@@ -74,7 +75,7 @@ It has been integrated into our establishing [brain modeling ecosystem](https://
 
 ## Installation
 
-``braintrace`` can run on Python 3.10+ installed on Linux, MacOS, and Windows. You can install ``braintrace`` via pip:
+``braintrace`` requires Python 3.11 or newer and supports Linux, macOS, and Windows. You can install ``braintrace`` via pip:
 
 ```bash
 pip install braintrace --upgrade
@@ -86,6 +87,47 @@ Alternatively, you can install `BrainX`, which bundles `braintrace` with other c
 pip install BrainX -U
 ```
 
+## Quickstart
+
+Mark the trainable operations with an ETP primitive (here, via `braintrace.nn`
+layers, which use them internally), then compile the model into a learner and
+drive a sequence. There is no scan to write and no special parameter class —
+weights stay plain `brainstate.ParamState`.
+
+```python
+import brainstate
+import jax.numpy as jnp
+import braintrace
+
+class GRUNet(brainstate.nn.Module):
+    def __init__(self, n_in, n_rec, n_out):
+        super().__init__()
+        self.rnn = braintrace.nn.GRUCell(n_in, n_rec)
+        self.out = braintrace.nn.Linear(n_rec, n_out)
+
+    def update(self, x):
+        return self.out(self.rnn(x))
+
+model = GRUNet(3, 8, 1)
+
+# One call: initialise states, compile the eligibility-trace graph,
+# return a ready learner.
+learner = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros((1, 3)), batch_size=1)
+
+def step_loss(x, y):
+    return jnp.mean((learner(x) - y) ** 2)
+
+xs = brainstate.random.randn(100, 1, 3)   # (time, batch, features)
+ys = brainstate.random.randn(100, 1, 1)
+
+# Drives the sequence and accumulates online gradients -- O(1) memory in T.
+grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
+```
+
+Swap `braintrace.D_RTRL` for `pp_prop`, `SnAp`, `UORO`, `ThreeFactor`, `DNI`,
+`EProp`, the OSTL variants, or an explicit `braintrace.ETraceConfig` to change
+the learning rule; nothing else in the snippet changes.
+
 ## Documentation
 
 The official documentation is hosted on Read the Docs: [https://brainx.chaobrain.com/braintrace](https://brainx.chaobrain.com/braintrace)

{braintrace-0.2.3 → braintrace-0.2.5}/README.md

@@ -19,7 +19,7 @@ It has been integrated into our establishing [brain modeling ecosystem](https://
 
 ## Installation
 
-``braintrace`` can run on Python 3.10+ installed on Linux, MacOS, and Windows. You can install ``braintrace`` via pip:
+``braintrace`` requires Python 3.11 or newer and supports Linux, macOS, and Windows. You can install ``braintrace`` via pip:
 
 ```bash
 pip install braintrace --upgrade
@@ -31,6 +31,47 @@ Alternatively, you can install `BrainX`, which bundles `braintrace` with other c
 pip install BrainX -U
 ```
 
+## Quickstart
+
+Mark the trainable operations with an ETP primitive (here, via `braintrace.nn`
+layers, which use them internally), then compile the model into a learner and
+drive a sequence. There is no scan to write and no special parameter class —
+weights stay plain `brainstate.ParamState`.
+
+```python
+import brainstate
+import jax.numpy as jnp
+import braintrace
+
+class GRUNet(brainstate.nn.Module):
+    def __init__(self, n_in, n_rec, n_out):
+        super().__init__()
+        self.rnn = braintrace.nn.GRUCell(n_in, n_rec)
+        self.out = braintrace.nn.Linear(n_rec, n_out)
+
+    def update(self, x):
+        return self.out(self.rnn(x))
+
+model = GRUNet(3, 8, 1)
+
+# One call: initialise states, compile the eligibility-trace graph,
+# return a ready learner.
+learner = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros((1, 3)), batch_size=1)
+
+def step_loss(x, y):
+    return jnp.mean((learner(x) - y) ** 2)
+
+xs = brainstate.random.randn(100, 1, 3)   # (time, batch, features)
+ys = brainstate.random.randn(100, 1, 1)
+
+# Drives the sequence and accumulates online gradients -- O(1) memory in T.
+grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
+```
+
+Swap `braintrace.D_RTRL` for `pp_prop`, `SnAp`, `UORO`, `ThreeFactor`, `DNI`,
+`EProp`, the OSTL variants, or an explicit `braintrace.ETraceConfig` to change
+the learning rule; nothing else in the snippet changes.
+
 ## Documentation
 
 The official documentation is hosted on Read the Docs: [https://brainx.chaobrain.com/braintrace](https://brainx.chaobrain.com/braintrace)

{braintrace-0.2.3 → braintrace-0.2.5}/braintrace/__init__.py

@@ -41,8 +41,10 @@ downward:
    hidden->weight / hidden->hidden Jacobian computations.
 4. **Algorithms** — online-learning orchestrators: the exact algorithms
    :class:`D_RTRL` / :func:`pp_prop` / :class:`ES_D_RTRL`, and the SNN family
-   :class:`EProp`, :class:`OSTLRecurrent`, :class:`OSTLFeedforward`,
-   :class:`OTPE`, :class:`OTTT`, :class:`OSTTP`.
+   :class:`EProp`, :class:`OSTLRecurrent`, :class:`OSTLFeedforward`. Every one
+   of them carries the two sequence drivers of
+   :class:`SequenceDriverMixin` — ``etrace_grad`` and ``etrace_evolve`` — so a
+   caller never writes the per-step gradient-accumulation loop by hand.
 
 The :mod:`braintrace.nn` subpackage provides ready-made ETP-wired layers
 (linear maps, convolutions, recurrent cells, read-outs).
@@ -50,7 +52,13 @@ The :mod:`braintrace.nn` subpackage provides ready-made ETP-wired layers
 Notes
 -----
 The convenience entry point :func:`compile` wraps a model together with an
-algorithm into a single trainable object and is the recommended starting point.
+algorithm into a single trainable object and is the recommended starting point:
+one call replaces ``init_all_states`` + algorithm construction +
+``compile_graph`` (+ the vmap wrapper, with ``vmap=True``). The object it
+returns drives a sequence with :meth:`~SequenceDriverMixin.etrace_grad`
+(accumulate online gradients under a loss) or
+:meth:`~SequenceDriverMixin.etrace_evolve` (advance hidden state and the
+eligibility trace with no loss).
 The ``braintrace.MatMulOp`` / ``ETraceParam`` style names from the v0.1.x API
 are deprecated shims served lazily with a :class:`DeprecationWarning`; new code
 should mark parameters by routing them through ETP ops instead.
@@ -73,9 +81,12 @@ from typing import TYPE_CHECKING, Any
 from . import nn
 from ._compile import compile
 from ._algorithm import (
+    ETraceConfig,
     ETraceAlgorithm,
     EligibilityTrace,
     ETraceGraphExecutor,
+    ETraceVmap,
+    SequenceDriverMixin,
     ETraceVjpAlgorithm,
     ETraceVjpGraphExecutor,
     ParamDimVjpAlgorithm,
@@ -83,17 +94,21 @@ from ._algorithm import (
     pp_prop,
     ES_D_RTRL,
     IODimVjpAlgorithm,
+    SnAp,
+    RandomProjectionVjpAlgorithm,
+    UORO,
+    ThreeFactor,
+    DNI,
+    SyntheticGradient,
+    train_synthetic_gradient,
     EProp,
     OSTLRecurrent,
     OSTLFeedforward,
-    OTPE,
-    OTTT,
-    OSTTP,
     FixedRandomFeedback,
     KappaFilter,
-    PresynapticTrace,
 )
 from ._compiler import (
+    ControlFlowPolicy,
     ETraceGraph,
     compile_etrace_graph,
     HiddenParamOpRelation,
@@ -115,6 +130,9 @@ from ._compiler import (
 from ._op import (
     ETPPrimitive,
     matmul,
+    grouped_matmul,
+    embedding,
+    einsum,
     element_wise,
     conv,
     sparse_matmul,
@@ -152,9 +170,14 @@ __all__ = [
     '__version__',
     '__version_info__',
 
+    # learning-rule axes
+    'ETraceConfig',
+
     # algorithms
     'ETraceAlgorithm',
     'EligibilityTrace',
+    'ETraceVmap',
+    'SequenceDriverMixin',
     'ETraceVjpAlgorithm',
     'ETraceVjpGraphExecutor',
     'ParamDimVjpAlgorithm',
@@ -162,12 +185,22 @@ __all__ = [
     'pp_prop',
     'ES_D_RTRL',
     'IODimVjpAlgorithm',
+    'SnAp',
+    'RandomProjectionVjpAlgorithm',
+    'UORO',
+    'ThreeFactor',
+    'DNI',
+    'SyntheticGradient',
+    'train_synthetic_gradient',
 
     # one-call entry point
     'compile',
 
     # ETP primitives (user API)
     'matmul',
+    'grouped_matmul',
+    'embedding',
+    'einsum',
     'element_wise',
     'conv',
     'sparse_matmul',
@@ -185,6 +218,7 @@ __all__ = [
     'ETraceGraphExecutor',
 
     # compiler
+    'ControlFlowPolicy',
     'ETraceGraph',
     'compile_etrace_graph',
     'HiddenGroup',
@@ -212,12 +246,8 @@ __all__ = [
     'EProp',
     'OSTLRecurrent',
     'OSTLFeedforward',
-    'OTPE',
-    'OTTT',
-    'OSTTP',
     'FixedRandomFeedback',
     'KappaFilter',
-    'PresynapticTrace',
 
     # errors
     'NotSupportedError',

{braintrace-0.2.3 → braintrace-0.2.5}/braintrace/__init___test.py

@@ -33,7 +33,7 @@ import pytest
 
 import braintrace
 import braintrace._legacy as legacy
-from braintrace._algorithm.oracle_models import tanh_rnn
+from braintrace._testing.oracle_models import tanh_rnn
 
 
 # ===========================================================================
@@ -71,6 +71,10 @@ def _algo_constructors():
         'EProp': lambda m: braintrace.EProp(m),
         'OSTLRecurrent': lambda m: braintrace.OSTLRecurrent(m),
         'OSTLFeedforward': lambda m: braintrace.OSTLFeedforward(m, decay_or_rank=0.9),
+        'SnAp': lambda m: braintrace.SnAp(m, n=1),
+        'UORO': lambda m: braintrace.UORO(m),
+        # `modulatory` has no fallback: a standing modulator is required.
+        'ThreeFactor': lambda m: braintrace.ThreeFactor(m, modulator=0.5),
     }
 
 
@@ -87,33 +91,30 @@ def test_algorithm_is_usable_end_to_end(name):
     assert bool(jnp.all(jnp.isfinite(y)))
 
 
-# --- Task 3: SNN algos requiring leak + OSTTP B_list contract ----------------
-
-def test_ottt_otpe_require_explicit_leak():
-    model = tanh_rnn(seed=0).factory()
+def test_dni_is_usable_end_to_end():
+    # DNI is the one preset whose second collaborator is sized from the compiled
+    # graph, so it cannot join the shared map above: build, compile, then attach.
+    model = tanh_rnn(n_in=3, n_rec=4, seed=0).factory()
     brainstate.nn.init_all_states(model, batch_size=1)
-    with pytest.raises(TypeError):
-        braintrace.OTTT(model)          # leak is keyword-only & required
-    with pytest.raises(TypeError):
-        braintrace.OTPE(model)
-
-
-def test_ottt_otpe_construct_with_leak():
-    for ctor in (lambda m: braintrace.OTTT(m, leak=0.9),
-                 lambda m: braintrace.OTPE(m, leak=0.9)):
-        model = tanh_rnn(seed=0).factory()
-        brainstate.nn.init_all_states(model, batch_size=1)
-        algo = ctor(model)
-        assert algo is not None
+    algo = braintrace.DNI(model)
+    x0 = jnp.ones((3,), dtype='float32')
+    algo.compile_graph(x0)
+    algo.init_etrace_state()
+    algo.attach_synthesizer(
+        braintrace.SyntheticGradient(algo.group_signal_shapes()))
+    y = algo(x0)
+    assert y.shape == (1, 4)
+    assert bool(jnp.all(jnp.isfinite(y)))
 
 
-def test_osttp_constructs_with_B_list():
-    model = tanh_rnn(n_in=3, n_rec=4, seed=0).factory()
-    brainstate.nn.init_all_states(model, batch_size=1)
-    # OSTTP needs one random-feedback matrix per HiddenGroup; B.shape[1] == n_l (=4).
-    B_list = [jnp.eye(4, dtype='float32')]
-    algo = braintrace.OSTTP(model, B_list)
-    assert algo is not None
+@pytest.mark.parametrize('name', ['UORO', 'ThreeFactor', 'DNI',
+                                  'SyntheticGradient', 'train_synthetic_gradient',
+                                  'RandomProjectionVjpAlgorithm'])
+def test_p4_symbols_are_exported_from_both_all_lists(name):
+    from braintrace import _algorithm
+    assert name in braintrace.__all__, f'{name} missing from braintrace.__all__'
+    assert name in _algorithm.__all__, f'{name} missing from _algorithm.__all__'
+    assert getattr(braintrace, name) is getattr(_algorithm, name)
 
 
 # --- Task 4: dynamic weight assignment raises NotSupportedError --------------
@@ -148,7 +149,14 @@ def test_weight_used_inside_scan_raises_not_implemented():
     within a control-flow op. Previously (F-SCAN-WEIGHT) the message construction
     indexed the hidden-path map with the weight var and a KeyError escaped instead;
     the message now renders the weight variable directly, so the intended
-    NotImplementedError surfaces."""
+    NotImplementedError surfaces.
+
+    Since Phase 2 canonicalization, statically short scans are unrolled at
+    extraction time instead of erroring, so this uses a length above
+    ``ControlFlowPolicy.scan_unroll_limit`` to reach the unsupported-op path
+    (after a SCAN_UNROLL_SKIPPED warning). Since Phase 4, structured scan
+    descent handles over-limit ETP scans by default, so this pins the
+    opt-out path via ``ControlFlowPolicy(scan_descent='off')``."""
 
     class ScanModel(brainstate.nn.Module):
         def __init__(self):
@@ -160,15 +168,17 @@ def test_weight_used_inside_scan_raises_not_implemented():
         def update(self, x):
             def body(carry, _):
                 return braintrace.matmul(carry, self.w.value), None  # weight inside scan
-            out, _ = jax.lax.scan(body, self.h.value, xs=None, length=2)
+            out, _ = jax.lax.scan(body, self.h.value, xs=None, length=64)
             self.h.value = jax.nn.tanh(out)
             return self.h.value
 
     model = ScanModel()
     brainstate.nn.init_all_states(model, batch_size=1)
-    algo = braintrace.D_RTRL(model)
+    algo = braintrace.D_RTRL(
+        model, control_flow=braintrace.ControlFlowPolicy(scan_descent='off'))
     with pytest.raises(NotImplementedError):
-        algo.compile_graph(jnp.ones((1, 4), dtype='float32'))
+        with pytest.warns(UserWarning, match='NOT unrolled'):
+            algo.compile_graph(jnp.ones((1, 4), dtype='float32'))
 
 
 # --- Task 6: legacy + nn deprecations emit DeprecationWarning ----------------
@@ -191,6 +201,14 @@ def test_nn_forwarded_name_warns():
         _ = nn.LayerNorm
 
 
+def test_nn_forwarded_state_name_warns_and_resolves():
+    import brainpy.state
+    import braintrace.nn as nn
+    with pytest.warns(DeprecationWarning):
+        forwarded = nn.LIF
+    assert forwarded is brainpy.state.LIF
+
+
 def test_nn_unknown_name_raises_attribute_error():
     import braintrace.nn as nn
     with pytest.raises(AttributeError):
@@ -217,8 +235,20 @@ def test_iodim_lives_in_io_dim_vjp_module():
 def test_expected_rnn_cells_exist():
     import braintrace.nn as nn
     for cell in ('ValinaRNNCell', 'GRUCell', 'MGUCell', 'LSTMCell', 'URLSTMCell',
-                 'MinimalRNNCell', 'MiniGRU', 'MiniLSTM', 'LRUCell'):
+                 'MinimalRNNCell', 'MiniGRU', 'MiniLSTM', 'LRUCell', 'CFNCell'):
         assert hasattr(nn, cell), f"missing cell: {cell}"
+        assert cell in nn.__all__, f"cell not in __all__: {cell}"
+
+
+def test_nn_dir_advertises_deprecated_forwards():
+    # ``braintrace.nn.__getattr__`` forwards ~50 names to brainstate.nn /
+    # brainpy.state. Without ``__dir__`` none of them is discoverable, so the
+    # forwarding is invisible to tab-completion and to ``dir()``-based tooling.
+    import braintrace.nn as nn
+    names = dir(nn)
+    assert 'ReLU' in names  # forwarded to brainstate.nn
+    assert 'LIF' in names  # forwarded to brainpy.state
+    assert set(nn.__all__).issubset(names)
 
 
 # ===========================================================================

{braintrace-0.2.3 → braintrace-0.2.5}/braintrace/_algorithm/__init__.py

@@ -17,32 +17,40 @@
 
 Groups the core ETrace infrastructure (``ETraceAlgorithm``, ``EligibilityTrace``,
 ``ETraceGraphExecutor``), VJP-based algorithms (D-RTRL, pp_prop / ES-D-RTRL),
-and paper-faithful SNN algorithms (EProp, OSTLRecurrent, OSTLFeedforward, OTPE,
-OTTT, OSTTP).
+and paper-faithful SNN algorithms (EProp, OSTLRecurrent, OSTLFeedforward).
 """
 
 from __future__ import annotations
 
-from ._common import FixedRandomFeedback, KappaFilter, PresynapticTrace
+from ._common import FixedRandomFeedback, KappaFilter
+from .axes import ETraceConfig
 from .base import ETraceAlgorithm, EligibilityTrace
 from .d_rtrl import D_RTRL
 from .e_prop import EProp
 from .graph_executor import ETraceGraphExecutor
 from .ostl import OSTLFeedforward, OSTLRecurrent
-from .osttp import OSTTP
-from .otpe import OTPE
-from .ottt import OTTT
 from .io_dim_vjp import IODimVjpAlgorithm
 from .param_dim_vjp import ParamDimVjpAlgorithm
+from .random_projection_vjp import RandomProjectionVjpAlgorithm
 from .pp_prop import ES_D_RTRL, pp_prop  # ES_D_RTRL: back-compat alias
+from .sequence import ETraceVmap, SequenceDriverMixin
+from .snap_n import SnAp
+from .dni import DNI, SyntheticGradient, train_synthetic_gradient
+from .three_factor import ThreeFactor
+from .uoro import UORO
 from .vjp_base import ETraceVjpAlgorithm
 from .vjp_graph_executor import ETraceVjpGraphExecutor
 
 __all__ = [
+    # axes
+    'ETraceConfig',
     # core
     'ETraceAlgorithm',
     'EligibilityTrace',
     'ETraceGraphExecutor',
+    # sequence drivers
+    'ETraceVmap',
+    'SequenceDriverMixin',
     # VJP
     'ETraceVjpAlgorithm',
     'ETraceVjpGraphExecutor',
@@ -51,14 +59,17 @@ __all__ = [
     'pp_prop',
     'ES_D_RTRL',
     'IODimVjpAlgorithm',
+    'SnAp',
+    'RandomProjectionVjpAlgorithm',
+    'UORO',
+    'ThreeFactor',
+    'DNI',
+    'SyntheticGradient',
+    'train_synthetic_gradient',
     # SNN
     'EProp',
     'OSTLRecurrent',
     'OSTLFeedforward',
-    'OTPE',
-    'OTTT',
-    'OSTTP',
     'FixedRandomFeedback',
     'KappaFilter',
-    'PresynapticTrace',
 ]