braintrace 0.2.1__tar.gz → 0.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.1
+Version: 0.2.3
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -26,10 +26,10 @@ 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: brainevent
+Requires-Dist: brainpy.state
 Requires-Dist: braintools
 Provides-Extra: cpu
 Requires-Dist: jax[cpu]; extra == "cpu"
@@ -41,10 +41,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 +66,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.1 → braintrace-0.2.3}/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.1 → braintrace-0.2.3}/braintrace/__init__.py

@@ -66,11 +66,13 @@ Examples
 """
 
 
-from typing import TYPE_CHECKING
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from . import nn
 from ._compile import compile
-from ._etrace_algorithms import (
+from ._algorithm import (
     ETraceAlgorithm,
     EligibilityTrace,
     ETraceGraphExecutor,
@@ -91,7 +93,7 @@ from ._etrace_algorithms import (
     KappaFilter,
     PresynapticTrace,
 )
-from ._etrace_compiler import (
+from ._compiler import (
     ETraceGraph,
     compile_etrace_graph,
     HiddenParamOpRelation,
@@ -106,10 +108,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 +201,7 @@ __all__ = [
 
     # compiler diagnostics
     'CompilationRecord',
+    'CompilationReport',
     'DiagnosticKind',
     'DiagnosticLevel',
 
@@ -242,7 +246,7 @@ _DEPRECATED_LEGACY = {
 }
 
 
-def __getattr__(name):
+def __getattr__(name: str) -> Any:
     if name in _DEPRECATED_LEGACY:
         import warnings
         warnings.warn(
@@ -256,5 +260,5 @@ def __getattr__(name):
     raise AttributeError(f'module {__name__!r} has no attribute {name!r}')
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return sorted(list(__all__) + list(_DEPRECATED_LEGACY))

braintrace-0.2.1/braintrace/api_contract_test.py → braintrace-0.2.3/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.1/braintrace/_etrace_algorithms → braintrace-0.2.3/braintrace/_algorithm}/__init__.py

@@ -21,6 +21,8 @@ and paper-faithful SNN algorithms (EProp, OSTLRecurrent, OSTLFeedforward, OTPE,
 OTTT, OSTTP).
 """
 
+from __future__ import annotations
+
 from ._common import FixedRandomFeedback, KappaFilter, PresynapticTrace
 from .base import ETraceAlgorithm, EligibilityTrace
 from .d_rtrl import D_RTRL

{braintrace-0.2.1/braintrace/_etrace_algorithms → braintrace-0.2.3/braintrace/_algorithm}/_common.py

@@ -15,15 +15,17 @@
 
 """Shared helpers for the SNN online-learning algorithms."""
 
+from __future__ import annotations
+
 from functools import partial
 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
+from braintrace._typing import ArrayLike, PyTree
 
 __all__ = [
     'PresynapticTrace',
@@ -42,12 +44,12 @@ class _ZeroResetState(brainstate.ShortTermState):
     for ``batch_size`` (or prepended for scalar-shaped states).
     """
 
-    def __init__(self, init_value):
+    def __init__(self, init_value: Any) -> None:
         super().__init__(init_value)
         self._init_shape = jnp.shape(init_value)
         self._init_dtype = init_value.dtype
 
-    def reset_state(self, batch_size: Optional[int] = None, **kwargs):
+    def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None:
         """Re-zero the state at its original shape.
 
         Parameters
@@ -105,13 +107,13 @@ class PresynapticTrace(_ZeroResetState):
 
     __module__ = 'braintrace'
 
-    def __init__(self, init_value, leak: float):
+    def __init__(self, init_value: Any, leak: float) -> None:
         super().__init__(init_value)
         if not (0.0 < leak < 1.0):
             raise ValueError(f'leak must be in (0, 1); got {leak}')
         self.leak = float(leak)
 
-    def update(self, x):
+    def update(self, x: ArrayLike) -> Any:
         r"""Apply one accumulation step :math:`\hat{a} \leftarrow \lambda \cdot \hat{a} + x`.
 
         Parameters
@@ -163,13 +165,13 @@ class KappaFilter(_ZeroResetState):
 
     __module__ = 'braintrace'
 
-    def __init__(self, init_value, kappa: float):
+    def __init__(self, init_value: Any, kappa: float) -> None:
         super().__init__(init_value)
         if not (0.0 <= kappa < 1.0):
             raise ValueError(f'kappa must be in [0, 1); got {kappa}')
         self.kappa = float(kappa)
 
-    def update(self, x):
+    def update(self, x: ArrayLike) -> Any:
         r"""Apply one low-pass step :math:`x_{\mathrm{filt}} \leftarrow (1-\kappa) x + \kappa\, x_{\mathrm{filt}}`.
 
         Parameters
@@ -230,14 +232,14 @@ class FixedRandomFeedback:
 
     __module__ = 'braintrace'
 
-    def __init__(self, n_target: int, n_layer: int, key, init_scale: float = 0.1):
+    def __init__(self, n_target: int, n_layer: int, key: Any, init_scale: float = 0.1) -> None:
         self.B = jax.lax.stop_gradient(
             init_scale * jax.random.normal(key, (n_target, n_layer))
         )
         self.n_target = int(n_target)
         self.n_layer = int(n_layer)
 
-    def project(self, y_target):
+    def project(self, y_target: Any) -> Any:
         """Project the target onto the frozen feedback matrix.
 
         Parameters
@@ -267,7 +269,7 @@ def extract_y_target(args: tuple, *, index: int = -1) -> Optional[jax.Array]:
 def _reset_state_in_a_dict(
     state_dict: Dict[Any, brainstate.State],
     batch_size: Optional[int],
-):
+) -> None:
     """
     Reset the values in a dictionary of states to zero.
 
@@ -293,7 +295,7 @@ def _reset_state_in_a_dict(
 def _zeros_like_batch_or_not(
     batch_size: Optional[int],
     x: jax.Array
-):
+) -> Any:
     """
     Create a zeros array with the same shape and type as the input array,
     optionally including a batch dimension.
@@ -325,7 +327,7 @@ def _batched_zeros_like(
     batch_size: Optional[int],
     num_state: int,  # the number of hidden states
     x: jax.Array  # the input array
-):
+) -> Any:
     """
     Create a batched zeros array with the same shape as the input array,
     extended by the number of hidden states.
@@ -354,7 +356,7 @@ def _batched_zeros_like(
         return u.math.zeros((batch_size, *x.shape, num_state), x.dtype)
 
 
-def _sum_dim(xs: jax.Array, axis: int = -1):
+def _sum_dim(xs: jax.Array, axis: int = -1) -> Any:
     """
     Sums the elements along the last dimension of each array in a PyTree.
 
@@ -373,7 +375,7 @@ def _sum_dim(xs: jax.Array, axis: int = -1):
     return jax.tree.map(lambda x: u.math.sum(x, axis=axis), xs)
 
 
-def _unit_safe_add(a, b):
+def _unit_safe_add(a: Any, b: Any) -> Any:
     """Add two leaves, stripping units only when one side has units and the other does not.
 
     Gradient contributions for the same weight may come from paths that
@@ -390,7 +392,7 @@ def _unit_safe_add(a, b):
     return u.math.add(a, b)
 
 
-def _extract_leaf(pytree_val: PyTree, leaf_idx: int):
+def _extract_leaf(pytree_val: PyTree, leaf_idx: int) -> Any:
     """Return the leaf at ``leaf_idx`` in ``jax.tree.leaves(pytree_val)``.
 
     Bare arrays (treedef with a single leaf) return the array unchanged.
@@ -409,7 +411,7 @@ def _extract_leaf(pytree_val: PyTree, leaf_idx: int):
 def _wrap_leaves_as_pytree(
     reference_pytree: PyTree,
     leaf_grads: Dict[int, jax.Array],
-):
+) -> Any:
     """Build a pytree matching ``reference_pytree`` with ``leaf_grads``
     inserted at the given leaf indices; any other leaf is zero-filled.
 
@@ -441,7 +443,7 @@ def _wrap_leaves_as_pytree(
 
 
 def _route_grads_by_path(
-    relation,
+    relation: Any,
     per_key_grads: Dict[str, jax.Array],
     weight_vals: Dict[Any, PyTree],
     target_dict: Dict[Any, PyTree],
@@ -478,7 +480,7 @@ def _update_dict(
     key: Any,
     value: PyTree,
     error_when_no_key: Optional[bool] = False
-):
+) -> None:
     """Update the dictionary.
 
     If the key exists, then add the value to the existing value.

{braintrace-0.2.1/braintrace/_etrace_algorithms → braintrace-0.2.3/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.1/braintrace/_etrace_algorithms → braintrace-0.2.3/braintrace/_algorithm}/base.py

@@ -19,11 +19,13 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 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
 
@@ -89,7 +91,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         model: brainstate.nn.Module,
         graph_executor: ETraceGraphExecutor,
         name: Optional[str] = None,
-    ):
+    ) -> None:
         super().__init__(name=name)  # type: ignore[call-arg]  # brainstate hides Module.__init__ from type checkers
 
         # the model
@@ -131,6 +133,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:
         """
@@ -188,7 +212,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         assert self._other_states is not None
         return self._other_states
 
-    def _split_state(self):
+    def _split_state(self) -> None:
         # --- the state separation --- #
         #
         # [NOTE]
@@ -206,7 +230,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
             self._other_states
         ) = self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...)
 
-    def compile_graph(self, *args) -> None:
+    def compile_graph(self, *args: Any) -> None:
         r"""
         Compile the eligibility trace graph of the relationship between etrace weights, states and operators.
 
@@ -222,6 +246,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
@@ -261,13 +297,30 @@ 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:
+    def __call__(self, *args: Any) -> Any:
         """
         Update the model and the eligibility trace states.
 
@@ -283,7 +336,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         return self.update(*args)
 
-    def update(self, *args) -> Any:
+    def update(self, *args: Any) -> Any:
         """
         Update the model and the eligibility trace states.
 
@@ -304,7 +357,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         raise NotImplementedError
 
-    def init_etrace_state(self, *args, **kwargs) -> None:
+    def init_etrace_state(self, *args: Any, **kwargs: Any) -> None:
         """
         Initialize the eligibility trace states of the etrace algorithm.
 

{braintrace-0.2.1/braintrace/_etrace_algorithms → braintrace-0.2.3/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()

{braintrace-0.2.1/braintrace/_etrace_algorithms → braintrace-0.2.3/braintrace/_algorithm}/d_rtrl.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 from .param_dim_vjp import ParamDimVjpAlgorithm
 
 __all__ = [