braintrace 0.2.0__tar.gz → 0.2.2__tar.gz

{braintrace-0.2.0 → braintrace-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.0
+Version: 0.2.2
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -26,10 +26,9 @@ Classifier: Topic :: Software Development :: Libraries
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: brainstate>=0.2.2
+Requires-Dist: brainstate>=0.5.2
 Requires-Dist: brainunit
-Requires-Dist: brainstate
-Requires-Dist: brainpy-state
+Requires-Dist: brainpy.state
 Requires-Dist: braintools
 Provides-Extra: cpu
 Requires-Dist: jax[cpu]; extra == "cpu"
@@ -41,10 +40,12 @@ Provides-Extra: tpu
 Requires-Dist: jax[tpu]; extra == "tpu"
 Provides-Extra: testing
 Requires-Dist: pytest; extra == "testing"
+Requires-Dist: pytest-cov; extra == "testing"
 Requires-Dist: jax[cpu]; extra == "testing"
 Requires-Dist: hypothesis; extra == "testing"
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
 Requires-Dist: jax[cpu]; extra == "dev"
 Requires-Dist: hypothesis; extra == "dev"
 Requires-Dist: mypy>=1.8; extra == "dev"
@@ -64,6 +65,7 @@ Dynamic: license-file
   	<a href="https://brainx.chaobrain.com/braintrace/"><img alt="Documentation" src="https://readthedocs.org/projects/braintrace/badge/?version=latest"></a>
   	<a href="https://badge.fury.io/py/braintrace"><img alt="PyPI version" src="https://badge.fury.io/py/braintrace.svg"></a>
     <a href="https://github.com/chaobrain/braintrace/actions/workflows/CI.yml"><img alt="Continuous Integration" src="https://github.com/chaobrain/braintrace/actions/workflows/CI.yml/badge.svg"></a>
+    <a href="https://codecov.io/gh/chaobrain/braintrace"><img alt="Code Coverage" src="https://codecov.io/gh/chaobrain/braintrace/branch/main/graph/badge.svg"></a>
 </p>
 
 [``braintrace``](https://github.com/chaobrain/braintrace) provides online learning algorithms for biological neural networks.

{braintrace-0.2.0 → braintrace-0.2.2}/README.md

@@ -11,6 +11,7 @@
   	<a href="https://brainx.chaobrain.com/braintrace/"><img alt="Documentation" src="https://readthedocs.org/projects/braintrace/badge/?version=latest"></a>
   	<a href="https://badge.fury.io/py/braintrace"><img alt="PyPI version" src="https://badge.fury.io/py/braintrace.svg"></a>
     <a href="https://github.com/chaobrain/braintrace/actions/workflows/CI.yml"><img alt="Continuous Integration" src="https://github.com/chaobrain/braintrace/actions/workflows/CI.yml/badge.svg"></a>
+    <a href="https://codecov.io/gh/chaobrain/braintrace"><img alt="Code Coverage" src="https://codecov.io/gh/chaobrain/braintrace/branch/main/graph/badge.svg"></a>
 </p>
 
 [``braintrace``](https://github.com/chaobrain/braintrace) provides online learning algorithms for biological neural networks.

{braintrace-0.2.0 → braintrace-0.2.2}/braintrace/__init__.py

@@ -70,7 +70,7 @@ from typing import TYPE_CHECKING
 
 from . import nn
 from ._compile import compile
-from ._etrace_algorithms import (
+from ._algorithm import (
     ETraceAlgorithm,
     EligibilityTrace,
     ETraceGraphExecutor,
@@ -91,7 +91,7 @@ from ._etrace_algorithms import (
     KappaFilter,
     PresynapticTrace,
 )
-from ._etrace_compiler import (
+from ._compiler import (
     ETraceGraph,
     compile_etrace_graph,
     HiddenParamOpRelation,
@@ -106,10 +106,11 @@ from ._etrace_compiler import (
     ModuleInfo,
     extract_module_info,
     CompilationRecord,
+    CompilationReport,
     DiagnosticKind,
     DiagnosticLevel,
 )
-from ._etrace_op import (
+from ._op import (
     ETPPrimitive,
     matmul,
     element_wise,
@@ -198,6 +199,7 @@ __all__ = [
 
     # compiler diagnostics
     'CompilationRecord',
+    'CompilationReport',
     'DiagnosticKind',
     'DiagnosticLevel',
 

braintrace-0.2.0/braintrace/api_contract_test.py → braintrace-0.2.2/braintrace/__init___test.py

@@ -13,10 +13,18 @@
 # limitations under the License.
 # ==============================================================================
 
-"""L0 public-API contract: every ``braintrace.__all__`` symbol resolves to a
-usable object; public algorithms run a minimal end-to-end step; documented error
-paths fire on their real triggers; legacy/nn deprecations warn. Also pins the
-implementation facts that drifted from CLAUDE.md (F-17)."""
+"""Tests for the ``braintrace`` package root (``braintrace/__init__.py``).
+
+Two concerns, both owned by the package root:
+
+* **Public-API contract** — every ``braintrace.__all__`` symbol resolves to a
+  usable object; public algorithms run a minimal end-to-end step; documented
+  error paths fire on their real triggers; legacy/nn deprecations warn; and the
+  implementation facts that drifted from CLAUDE.md (F-17) stay pinned.
+* **Legacy deprecation forwarding** — the v0.1.x shims served lazily by
+  ``__getattr__`` warn on access, stay out of ``__all__``, and appear in
+  ``__dir__``.
+"""
 
 import brainstate
 import jax
@@ -24,8 +32,13 @@ import jax.numpy as jnp
 import pytest
 
 import braintrace
-from braintrace._etrace_algorithms.oracle_models import tanh_rnn
+import braintrace._legacy as legacy
+from braintrace._algorithm.oracle_models import tanh_rnn
+
 
+# ===========================================================================
+# Public-API contract
+# ===========================================================================
 
 # --- Task 1: every __all__ symbol resolves to a non-None object --------------
 
@@ -130,18 +143,12 @@ def test_dynamic_weight_assignment_raises_not_supported():
 
 # --- Task 5: weight inside control flow raises NotImplementedError -----------
 
-@pytest.mark.xfail(
-    strict=True,
-    reason="F-SCAN-WEIGHT: the weight-in-control-flow guard intends to raise "
-           "NotImplementedError, but its error-message construction at "
-           "base.py:133 indexes invar_to_hidden_path with the *weight* var "
-           "(absent from that hidden-path map), so a KeyError escapes instead. "
-           "The guard fires, but the wrong exception type surfaces.",
-)
 def test_weight_used_inside_scan_raises_not_implemented():
-    """check_unsupported_op should raise NotImplementedError when a weight var is
-    used within a control-flow op. Today it raises KeyError while building the
-    message (F-SCAN-WEIGHT) — this test pins the intended contract via xfail."""
+    """check_unsupported_op raises NotImplementedError when a weight var is used
+    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."""
 
     class ScanModel(brainstate.nn.Module):
         def __init__(self):
@@ -193,8 +200,8 @@ def test_nn_unknown_name_raises_attribute_error():
 # --- Task 7: lock the drifted implementation facts (F-17) --------------------
 
 def test_ostl_is_two_classes_not_a_factory():
-    from braintrace._etrace_algorithms.param_dim_vjp import ParamDimVjpAlgorithm
-    from braintrace._etrace_algorithms.io_dim_vjp import IODimVjpAlgorithm
+    from braintrace._algorithm.param_dim_vjp import ParamDimVjpAlgorithm
+    from braintrace._algorithm.io_dim_vjp import IODimVjpAlgorithm
     assert isinstance(braintrace.OSTLRecurrent, type)
     assert isinstance(braintrace.OSTLFeedforward, type)
     assert issubclass(braintrace.OSTLRecurrent, ParamDimVjpAlgorithm)
@@ -202,7 +209,7 @@ def test_ostl_is_two_classes_not_a_factory():
 
 
 def test_iodim_lives_in_io_dim_vjp_module():
-    from braintrace._etrace_algorithms.io_dim_vjp import IODimVjpAlgorithm
+    from braintrace._algorithm.io_dim_vjp import IODimVjpAlgorithm
     assert braintrace.IODimVjpAlgorithm is IODimVjpAlgorithm
     assert braintrace.pp_prop is braintrace.ES_D_RTRL  # aliases
 
@@ -212,3 +219,42 @@ def test_expected_rnn_cells_exist():
     for cell in ('ValinaRNNCell', 'GRUCell', 'MGUCell', 'LSTMCell', 'URLSTMCell',
                  'MinimalRNNCell', 'MiniGRU', 'MiniLSTM', 'LRUCell'):
         assert hasattr(nn, cell), f"missing cell: {cell}"
+
+
+# ===========================================================================
+# Legacy v0.1.x deprecation forwarding
+# ===========================================================================
+
+_LEGACY_NAMES = [
+    'ETraceOp', 'MatMulOp', 'ElemWiseOp', 'ConvOp', 'SpMatMulOp', 'LoraOp',
+    'ETraceParam', 'ElemWiseParam', 'NonTempParam',
+    'FakeETraceParam', 'FakeElemWiseParam',
+]
+
+
+@pytest.mark.parametrize('name', _LEGACY_NAMES)
+def test_legacy_access_warns_and_returns_class(name):
+    with pytest.warns(DeprecationWarning):
+        obj = getattr(braintrace, name)
+    assert obj is getattr(legacy, name)
+
+
+@pytest.mark.parametrize('name', _LEGACY_NAMES)
+def test_legacy_names_not_in_all(name):
+    assert name not in braintrace.__all__
+
+
+def test_from_import_warns():
+    with pytest.warns(DeprecationWarning):
+        from braintrace import MatMulOp  # noqa: F401
+
+
+def test_unknown_attribute_raises_attribute_error():
+    with pytest.raises(AttributeError):
+        _ = braintrace.ThisNameDoesNotExist
+
+
+def test_legacy_names_in_dir():
+    d = dir(braintrace)
+    for name in _LEGACY_NAMES:
+        assert name in d

{braintrace-0.2.0/braintrace/_etrace_algorithms → braintrace-0.2.2/braintrace/_algorithm}/_common.py

@@ -21,7 +21,9 @@ from typing import Any, Dict, Optional
 import brainstate
 import jax
 import jax.numpy as jnp
-import saiunit as u
+import brainunit as u
+
+from braintrace._typing import PyTree
 
 __all__ = [
     'PresynapticTrace',
@@ -388,7 +390,7 @@ def _unit_safe_add(a, b):
     return u.math.add(a, b)
 
 
-def _extract_leaf(pytree_val: brainstate.typing.PyTree, leaf_idx: int):
+def _extract_leaf(pytree_val: PyTree, leaf_idx: int):
     """Return the leaf at ``leaf_idx`` in ``jax.tree.leaves(pytree_val)``.
 
     Bare arrays (treedef with a single leaf) return the array unchanged.
@@ -405,7 +407,7 @@ def _extract_leaf(pytree_val: brainstate.typing.PyTree, leaf_idx: int):
 
 
 def _wrap_leaves_as_pytree(
-    reference_pytree: brainstate.typing.PyTree,
+    reference_pytree: PyTree,
     leaf_grads: Dict[int, jax.Array],
 ):
     """Build a pytree matching ``reference_pytree`` with ``leaf_grads``
@@ -441,8 +443,8 @@ def _wrap_leaves_as_pytree(
 def _route_grads_by_path(
     relation,
     per_key_grads: Dict[str, jax.Array],
-    weight_vals: Dict[Any, brainstate.typing.PyTree],
-    target_dict: Dict[Any, brainstate.typing.PyTree],
+    weight_vals: Dict[Any, PyTree],
+    target_dict: Dict[Any, PyTree],
 ) -> None:
     """Route per-key gradients from a dict-API rule into per-path pytrees.
 
@@ -474,7 +476,7 @@ def _route_grads_by_path(
 def _update_dict(
     the_dict: Dict,
     key: Any,
-    value: brainstate.typing.PyTree,
+    value: PyTree,
     error_when_no_key: Optional[bool] = False
 ):
     """Update the dictionary.

{braintrace-0.2.0/braintrace/_etrace_algorithms → braintrace-0.2.2/braintrace/_algorithm}/_common_test.py

@@ -20,9 +20,9 @@ import jax
 import jax.numpy as jnp
 import numpy.testing as npt
 import pytest
-import saiunit as u
+import brainunit as u
 
-from braintrace._etrace_algorithms._common import (
+from braintrace._algorithm._common import (
     FixedRandomFeedback,
     KappaFilter,
     PresynapticTrace,

{braintrace-0.2.0/braintrace/_etrace_algorithms → braintrace-0.2.2/braintrace/_algorithm}/base.py

@@ -23,7 +23,7 @@ from typing import Dict, Any, Optional
 
 import brainstate
 
-from braintrace._etrace_compiler import ETraceGraph
+from braintrace._compiler import ETraceGraph, CompilationReport
 from .graph_executor import ETraceGraphExecutor
 from .._typing import Path
 
@@ -90,7 +90,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         graph_executor: ETraceGraphExecutor,
         name: Optional[str] = None,
     ):
-        super().__init__(name=name)
+        super().__init__(name=name)  # type: ignore[call-arg]  # brainstate hides Module.__init__ from type checkers
 
         # the model
         if not isinstance(model, brainstate.nn.Module):
@@ -115,9 +115,9 @@ class ETraceAlgorithm(brainstate.nn.Module):
         self.running_index = brainstate.LongTermState(0)
 
         # other states
-        self._param_states = None
-        self._hidden_states = None
-        self._other_states = None
+        self._param_states: Optional[brainstate.util.FlattedDict] = None
+        self._hidden_states: Optional[brainstate.util.FlattedDict] = None
+        self._other_states: Optional[brainstate.util.FlattedDict] = None
 
     @property
     def graph(self) -> ETraceGraph:
@@ -131,6 +131,28 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         return self.graph_executor.graph
 
+    @property
+    def report(self) -> CompilationReport:
+        """Structured, read-only report of the compiled eligibility-trace graph.
+
+        Returns
+        -------
+        CompilationReport
+            A view over :attr:`graph`.
+
+        Raises
+        ------
+        RuntimeError
+            If accessed before the graph has been compiled.
+        """
+        if not self.is_compiled:
+            raise RuntimeError(
+                'The eligibility-trace graph is not compiled yet. Build the '
+                'learner with braintrace.compile(...) (or call .compile_graph(...)) '
+                'before accessing .report.'
+            )
+        return CompilationReport(self.graph)
+
     @property
     def executor(self) -> ETraceGraphExecutor:
         """
@@ -144,7 +166,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         return self.graph_executor
 
     @property
-    def param_states(self) -> brainstate.util.FlattedDict[Path, brainstate.ParamState]:
+    def param_states(self) -> brainstate.util.FlattedDict:
         """
         Get the parameter weight states.
 
@@ -155,10 +177,11 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         if self._param_states is None:
             self._split_state()
+        assert self._param_states is not None
         return self._param_states
 
     @property
-    def hidden_states(self) -> brainstate.util.FlattedDict[Path, brainstate.HiddenState]:
+    def hidden_states(self) -> brainstate.util.FlattedDict:
         """
         Get the hidden states.
 
@@ -169,10 +192,11 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         if self._hidden_states is None:
             self._split_state()
+        assert self._hidden_states is not None
         return self._hidden_states
 
     @property
-    def other_states(self) -> brainstate.util.FlattedDict[Path, brainstate.State]:
+    def other_states(self) -> brainstate.util.FlattedDict:
         """
         Get the other states.
 
@@ -183,6 +207,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         if self._other_states is None:
             self._split_state()
+        assert self._other_states is not None
         return self._other_states
 
     def _split_state(self):
@@ -219,6 +244,18 @@ class ETraceAlgorithm(brainstate.nn.Module):
             The input arguments.
         """
 
+        # ``vmap_new_states`` / ``vmap2_new_states`` run an eager *discovery
+        # probe* that executes the surrounding ``init`` (including this call)
+        # once against throwaway, un-batched states before the real mapped
+        # pass. Compiling there would bind the executor to those probe states
+        # (which are discarded and left untagged), so the subsequent
+        # ``brainstate.nn.Vmap(..., vmap_states='new')`` would not cover them
+        # and writing a batched value raises ``BatchAxisError``. Defer to the
+        # real mapped pass, which creates the 'new'-tagged batched states.
+        _in_probe = getattr(brainstate.transform, 'in_new_state_probe', None)
+        if _in_probe is not None and _in_probe():
+            return
+
         if not self.is_compiled:
             # --- invalidate cached state splits --- #
             self._param_states = None
@@ -235,7 +272,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
             self.is_compiled = True
 
     @property
-    def path_to_states(self) -> brainstate.util.FlattedDict[Path, brainstate.State]:
+    def path_to_states(self) -> brainstate.util.FlattedDict:
         """
         Get the path to the states.
 
@@ -258,11 +295,28 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         return self.graph_executor.state_id_to_path
 
-    def show_graph(self) -> None:
-        """
-        Show the etrace graph.
+    def show_graph(
+        self,
+        verbose: bool = True,
+        return_msg: bool = False,
+    ) -> None | str:
+        """Display the eligibility-trace graph.
+
+        Delegates to :meth:`ETraceGraphExecutor.show_graph`.
+
+        Parameters
+        ----------
+        verbose : bool, optional
+            If True (default), print the summary to stdout.
+        return_msg : bool, optional
+            If True, return the summary string. Default False.
+
+        Returns
+        -------
+        None or str
+            The summary string if ``return_msg`` is True, else None.
         """
-        self.graph_executor.show_graph()
+        return self.graph_executor.show_graph(verbose=verbose, return_msg=return_msg)
 
     def __call__(self, *args) -> Any:
         """

{braintrace-0.2.0/braintrace/_etrace_algorithms → braintrace-0.2.2/braintrace/_algorithm}/base_test.py

@@ -21,8 +21,8 @@ import brainstate
 import jax.numpy as jnp
 
 import braintrace
-from braintrace._etrace_algorithms import ETraceAlgorithm, EligibilityTrace
-from braintrace._etrace_algorithms.graph_executor import ETraceGraphExecutor
+from braintrace._algorithm import ETraceAlgorithm, EligibilityTrace
+from braintrace._algorithm.graph_executor import ETraceGraphExecutor
 
 
 # ---------------------------------------------------------------------------
@@ -580,5 +580,28 @@ class TestETraceAlgorithmTwoLayerIntegration(unittest.TestCase):
         self.assertGreaterEqual(len(algo.hidden_states), 2)
 
 
+def test_report_before_compile_raises():
+    import brainstate, pytest, braintrace
+    from braintrace._algorithm.oracle_models import tanh_rnn
+    model = tanh_rnn(n_in=3, n_rec=4, seed=0).factory()
+    brainstate.nn.init_all_states(model, batch_size=1)
+    algo = braintrace.D_RTRL(model)
+    with pytest.raises(RuntimeError):
+        _ = algo.report
+
+
+def test_report_after_compile_matches_show_graph():
+    import brainstate, jax.numpy as jnp, braintrace
+    from braintrace import CompilationReport
+    from braintrace._algorithm.oracle_models import tanh_rnn
+    model = tanh_rnn(n_in=3, n_rec=4, seed=0).factory()
+    brainstate.nn.init_all_states(model, batch_size=1)
+    algo = braintrace.D_RTRL(model)
+    algo.compile_graph(jnp.ones((3,), 'float32'))
+    assert isinstance(algo.report, CompilationReport)
+    # show_graph(return_msg=True) is exactly report.to_str(1)
+    assert algo.show_graph(verbose=False, return_msg=True) == algo.report.to_str(1)
+
+
 if __name__ == '__main__':
     unittest.main()