braintrace 0.1.2__tar.gz → 0.2.1__tar.gz

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

@@ -1,19 +1,18 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.1.2
+Version: 0.2.1
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license
 Project-URL: Homepage, https://github.com/chaobrain/braintrace
 Project-URL: Bug Tracker, https://github.com/chaobrain/braintrace/issues
-Project-URL: Documentation, https://braintrace.readthedocs.io/
+Project-URL: Documentation, https://brainx.chaobrain.com/braintrace/
 Project-URL: Source Code, https://github.com/chaobrain/braintrace
 Keywords: computational neuroscience,brain-inspired computing,brain modeling,online learning
 Classifier: Natural Language :: English
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
@@ -24,7 +23,7 @@ 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.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: brainstate>=0.2.2
@@ -34,54 +33,41 @@ Requires-Dist: brainpy-state
 Requires-Dist: braintools
 Provides-Extra: cpu
 Requires-Dist: jax[cpu]; extra == "cpu"
-Requires-Dist: brainunit; extra == "cpu"
-Requires-Dist: brainstate; extra == "cpu"
-Requires-Dist: brainpy-state; extra == "cpu"
-Requires-Dist: braintools; extra == "cpu"
 Provides-Extra: cuda12
 Requires-Dist: jax[cuda12]; extra == "cuda12"
-Requires-Dist: brainunit; extra == "cuda12"
-Requires-Dist: brainstate; extra == "cuda12"
-Requires-Dist: brainpy-state; extra == "cuda12"
-Requires-Dist: braintools; extra == "cuda12"
 Provides-Extra: cuda13
 Requires-Dist: jax[cuda13]; extra == "cuda13"
-Requires-Dist: brainunit; extra == "cuda13"
-Requires-Dist: brainstate; extra == "cuda13"
-Requires-Dist: brainpy-state; extra == "cuda13"
-Requires-Dist: braintools; extra == "cuda13"
 Provides-Extra: tpu
 Requires-Dist: jax[tpu]; extra == "tpu"
-Requires-Dist: brainunit; extra == "tpu"
-Requires-Dist: brainstate; extra == "tpu"
-Requires-Dist: brainpy-state; extra == "tpu"
-Requires-Dist: braintools; extra == "tpu"
 Provides-Extra: testing
 Requires-Dist: pytest; extra == "testing"
 Requires-Dist: jax[cpu]; extra == "testing"
-Requires-Dist: brainunit; extra == "testing"
-Requires-Dist: brainstate; extra == "testing"
-Requires-Dist: brainpy-state; extra == "testing"
-Requires-Dist: braintools; extra == "testing"
+Requires-Dist: hypothesis; extra == "testing"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: jax[cpu]; extra == "dev"
+Requires-Dist: hypothesis; extra == "dev"
+Requires-Dist: mypy>=1.8; extra == "dev"
+Requires-Dist: build; extra == "dev"
 Dynamic: license-file
 
 <h1 align="center">BrainTrace</h1>
 <h2 align="center">Eligibility Trace-based Online Learning for Brain Dynamics</h2>
 
 <p align="center">
-  	<img alt="Header image of braintrace." src="https://raw.githubusercontent.com/chaobrain/braintrace/main/docs/_static/braintrace.png" width=40%>
+  	<img alt="Header image of braintrace." src="https://brainx.chaobrain.com/images/braintrace.webp" width=40%>
 </p>
 
 <p align="center">
 	<a href="https://pypi.org/project/braintrace/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/braintrace"></a>
 	<a href="https://github.com/chaobrain/braintrace/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
-  	<a href="https://braintrace.readthedocs.io/?badge=latest"><img alt="Documentation" src="https://readthedocs.org/projects/braintrace/badge/?version=latest"></a>
+  	<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>
 </p>
 
 [``braintrace``](https://github.com/chaobrain/braintrace) provides online learning algorithms for biological neural networks.
-It has been integrated into our establishing [brain modeling ecosystem](https://brainmodeling.readthedocs.io/).
+It has been integrated into our establishing [brain modeling ecosystem](https://brainx.chaobrain.com/).
 
 ## Installation
 
@@ -99,40 +85,38 @@ pip install BrainX -U
 
 ## Documentation
 
-The official documentation is hosted on Read the Docs: [https://braintrace.readthedocs.io](https://braintrace.readthedocs.io)
+The official documentation is hosted on Read the Docs: [https://brainx.chaobrain.com/braintrace](https://brainx.chaobrain.com/braintrace)
+
+## Citation
+
+If you use this package in your research, please cite:
+
+```bibtex
+
+@Article{Wang2026,
+  author={Wang, Chaoming
+          and Dong, Xingsi
+          and Ji, Zilong
+          and Xiao, Mingqing
+          and Jiang, Jiedong
+          and Liu, Xiao
+          and Huan, Yuxiang
+          and Wu, Si},
+  title={Model-agnostic linear-memory online learning in spiking neural networks},
+  journal={Nature Communications},
+  year={2026},
+  month={Jan},
+  day={19},
+  abstract={Spiking neural networks (SNNs) offer a promising paradigm for modeling brain dynamics and developing neuromorphic intelligence, yet an online learning system capable of training rich spiking dynamics over long horizons with low memory footprints has been missing. Existing online approaches either incur quadratic memory growth, sacrifice biological fidelity through oversimplified models, or lack end-to-end automated tooling. Here, we introduce BrainTrace, a model-agnostic, linear-memory, and automated online learning system for spiking neural networks. BrainTrace standardizes model specification to encompass diverse neuronal and synaptic dynamics; implements a linear-memory online learning rule by exploiting intrinsic properties of spiking dynamics; and provides a compiler that automatically generates optimized online-learning code for arbitrary user-defined models. Across diverse dynamics and tasks, BrainTrace achieves strong learning performance with a low memory footprint and high computational throughput. Critically, these properties enable online fitting of a whole-brain-scale Drosophila SNN that recapitulates region-level functional activity. By reconciling generality, efficiency, and usability, BrainTrace establishes a foundation for spiking network modeling at scale.},
+  issn={2041-1723},
+  doi={10.1038/s41467-026-68453-w},
+  url={https://doi.org/10.1038/s41467-026-68453-w},
+  publisher={Nature Publishing Group UK London}
+}
 
-[//]: # (## Citation)
-
-[//]: # ()
-[//]: # (If you use this package in your research, please cite:)
-
-[//]: # ()
-[//]: # (```bibtex)
-
-[//]: # (@article {Wang2024.09.24.614728,)
-
-[//]: # (	author = {Wang, Chaoming and Dong, Xingsi and Ji, Zilong and Jiang, Jiedong and Liu, Xiao and Wu, Si},)
-
-[//]: # (	title = {Enabling Scalable Online Learning in Spiking Neural Networks},)
-
-[//]: # (	elocation-id = {2024.09.24.614728},)
-
-[//]: # (	year = {2025},)
-
-[//]: # (	doi = {10.1101/2024.09.24.614728},)
-
-[//]: # (	publisher = {Cold Spring Harbor Laboratory},)
-
-[//]: # (	URL = {https://www.biorxiv.org/content/early/2025/07/27/2024.09.24.614728},)
-
-[//]: # (	eprint = {https://www.biorxiv.org/content/early/2025/07/27/2024.09.24.614728.full.pdf},)
-
-[//]: # (	journal = {bioRxiv})
-
-[//]: # (})
+```
 
-[//]: # (```)
 
 ## See also the ecosystem
 
-``braintrace`` is one part of our brain simulation ecosystem: https://brainmodeling.readthedocs.io/
+``braintrace`` is one part of our brain simulation ecosystem: https://brainx.chaobrain.com/

braintrace-0.2.1/README.md

@@ -0,0 +1,69 @@
+<h1 align="center">BrainTrace</h1>
+<h2 align="center">Eligibility Trace-based Online Learning for Brain Dynamics</h2>
+
+<p align="center">
+  	<img alt="Header image of braintrace." src="https://brainx.chaobrain.com/images/braintrace.webp" width=40%>
+</p>
+
+<p align="center">
+	<a href="https://pypi.org/project/braintrace/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/braintrace"></a>
+	<a href="https://github.com/chaobrain/braintrace/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
+  	<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>
+</p>
+
+[``braintrace``](https://github.com/chaobrain/braintrace) provides online learning algorithms for biological neural networks.
+It has been integrated into our establishing [brain modeling ecosystem](https://brainx.chaobrain.com/).
+
+## Installation
+
+``braintrace`` can run on Python 3.10+ installed on Linux, MacOS, and Windows. You can install ``braintrace`` via pip:
+
+```bash
+pip install braintrace --upgrade
+```
+
+Alternatively, you can install `BrainX`, which bundles `braintrace` with other compatible packages for a comprehensive brain modeling ecosystem:
+
+```bash
+pip install BrainX -U
+```
+
+## Documentation
+
+The official documentation is hosted on Read the Docs: [https://brainx.chaobrain.com/braintrace](https://brainx.chaobrain.com/braintrace)
+
+## Citation
+
+If you use this package in your research, please cite:
+
+```bibtex
+
+@Article{Wang2026,
+  author={Wang, Chaoming
+          and Dong, Xingsi
+          and Ji, Zilong
+          and Xiao, Mingqing
+          and Jiang, Jiedong
+          and Liu, Xiao
+          and Huan, Yuxiang
+          and Wu, Si},
+  title={Model-agnostic linear-memory online learning in spiking neural networks},
+  journal={Nature Communications},
+  year={2026},
+  month={Jan},
+  day={19},
+  abstract={Spiking neural networks (SNNs) offer a promising paradigm for modeling brain dynamics and developing neuromorphic intelligence, yet an online learning system capable of training rich spiking dynamics over long horizons with low memory footprints has been missing. Existing online approaches either incur quadratic memory growth, sacrifice biological fidelity through oversimplified models, or lack end-to-end automated tooling. Here, we introduce BrainTrace, a model-agnostic, linear-memory, and automated online learning system for spiking neural networks. BrainTrace standardizes model specification to encompass diverse neuronal and synaptic dynamics; implements a linear-memory online learning rule by exploiting intrinsic properties of spiking dynamics; and provides a compiler that automatically generates optimized online-learning code for arbitrary user-defined models. Across diverse dynamics and tasks, BrainTrace achieves strong learning performance with a low memory footprint and high computational throughput. Critically, these properties enable online fitting of a whole-brain-scale Drosophila SNN that recapitulates region-level functional activity. By reconciling generality, efficiency, and usability, BrainTrace establishes a foundation for spiking network modeling at scale.},
+  issn={2041-1723},
+  doi={10.1038/s41467-026-68453-w},
+  url={https://doi.org/10.1038/s41467-026-68453-w},
+  publisher={Nature Publishing Group UK London}
+}
+
+```
+
+
+## See also the ecosystem
+
+``braintrace`` is one part of our brain simulation ecosystem: https://brainx.chaobrain.com/

braintrace-0.2.1/braintrace/__init__.py

@@ -0,0 +1,260 @@
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+"""braintrace: online learning for recurrent networks via Eligibility Trace Propagation (ETP).
+
+``braintrace`` trains recurrent and spiking neural networks **online** — forward
+in time, without backpropagation through time (BPTT). Models mark their
+trainable operations with ETP user-API ops (for example :func:`matmul`,
+:func:`conv`, :func:`sparse_matmul`, :func:`lora_matmul`, :func:`element_wise`)
+instead of wrapping parameters in a special class. A compiler then walks the
+JAX ``jaxpr``, identifies those ETP primitives, and connects each parameter to
+the hidden states it influences so that eligibility traces can be propagated.
+
+The public API is organised in four layers, with dependencies pointing strictly
+downward:
+
+1. **ETP operators** — the user-facing ops (:func:`matmul`, :func:`conv`, ...),
+   the :class:`ETPPrimitive` class, and :func:`register_primitive` for adding
+   new ones.
+2. **Compiler** — :func:`compile_etrace_graph` and the analysis containers
+   (:class:`ETraceGraph`, :class:`ModuleInfo`, :class:`HiddenGroup`,
+   :class:`HiddenParamOpRelation`, :class:`HiddenPerturbation`) plus the
+   diagnostics types (:class:`CompilationRecord`, :class:`DiagnosticKind`,
+   :class:`DiagnosticLevel`).
+3. **Graph executor** — :class:`ETraceGraphExecutor` /
+   :class:`ETraceVjpGraphExecutor`, which run the forward pass and the
+   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`.
+
+The :mod:`braintrace.nn` subpackage provides ready-made ETP-wired layers
+(linear maps, convolutions, recurrent cells, read-outs).
+
+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.
+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.
+
+Examples
+--------
+.. code-block:: python
+
+    >>> import braintrace
+    >>> # the public API surface is enumerated by __all__
+    >>> 'matmul' in braintrace.__all__
+    True
+"""
+
+
+from typing import TYPE_CHECKING
+
+from . import nn
+from ._compile import compile
+from ._etrace_algorithms import (
+    ETraceAlgorithm,
+    EligibilityTrace,
+    ETraceGraphExecutor,
+    ETraceVjpAlgorithm,
+    ETraceVjpGraphExecutor,
+    ParamDimVjpAlgorithm,
+    D_RTRL,
+    pp_prop,
+    ES_D_RTRL,
+    IODimVjpAlgorithm,
+    EProp,
+    OSTLRecurrent,
+    OSTLFeedforward,
+    OTPE,
+    OTTT,
+    OSTTP,
+    FixedRandomFeedback,
+    KappaFilter,
+    PresynapticTrace,
+)
+from ._etrace_compiler import (
+    ETraceGraph,
+    compile_etrace_graph,
+    HiddenParamOpRelation,
+    find_hidden_param_op_relations_from_minfo,
+    find_hidden_param_op_relations_from_module,
+    HiddenGroup,
+    find_hidden_groups_from_minfo,
+    find_hidden_groups_from_module,
+    HiddenPerturbation,
+    add_hidden_perturbation_from_minfo,
+    add_hidden_perturbation_in_module,
+    ModuleInfo,
+    extract_module_info,
+    CompilationRecord,
+    DiagnosticKind,
+    DiagnosticLevel,
+)
+from ._etrace_op import (
+    ETPPrimitive,
+    matmul,
+    element_wise,
+    conv,
+    sparse_matmul,
+    lora_matmul,
+    register_primitive,
+)
+from ._grad_exponential import GradExpon
+from ._input_data import (
+    SingleStepData,
+    MultiStepData,
+)
+from ._misc import NotSupportedError, CompilationError
+from ._version import __version__, __version_info__
+
+if TYPE_CHECKING:
+    # The v0.1.x legacy shims are deprecated and served lazily via ``__getattr__``
+    # below. Re-import them here so static type checkers / IDEs can still resolve
+    # ``braintrace.MatMulOp`` etc.
+    from ._legacy import (
+        ConvOp,
+        ElemWiseOp,
+        ElemWiseParam,
+        ETraceOp,
+        ETraceParam,
+        FakeElemWiseParam,
+        FakeETraceParam,
+        LoraOp,
+        MatMulOp,
+        NonTempParam,
+        SpMatMulOp,
+    )
+
+__all__ = [
+    # version
+    '__version__',
+    '__version_info__',
+
+    # algorithms
+    'ETraceAlgorithm',
+    'EligibilityTrace',
+    'ETraceVjpAlgorithm',
+    'ETraceVjpGraphExecutor',
+    'ParamDimVjpAlgorithm',
+    'D_RTRL',
+    'pp_prop',
+    'ES_D_RTRL',
+    'IODimVjpAlgorithm',
+
+    # one-call entry point
+    'compile',
+
+    # ETP primitives (user API)
+    'matmul',
+    'element_wise',
+    'conv',
+    'sparse_matmul',
+    'lora_matmul',
+
+    # ETP primitive class & rule registration
+    'ETPPrimitive',
+    'register_primitive',
+
+    # input data
+    'SingleStepData',
+    'MultiStepData',
+
+    # graph executor
+    'ETraceGraphExecutor',
+
+    # compiler
+    'ETraceGraph',
+    'compile_etrace_graph',
+    'HiddenGroup',
+    'find_hidden_groups_from_minfo',
+    'find_hidden_groups_from_module',
+    'HiddenParamOpRelation',
+    'find_hidden_param_op_relations_from_minfo',
+    'find_hidden_param_op_relations_from_module',
+    'ModuleInfo',
+    'extract_module_info',
+    'HiddenPerturbation',
+    'add_hidden_perturbation_from_minfo',
+    'add_hidden_perturbation_in_module',
+
+    # compiler diagnostics
+    'CompilationRecord',
+    'DiagnosticKind',
+    'DiagnosticLevel',
+
+    # gradient utilities
+    'GradExpon',
+
+    # SNN online-learning algorithms
+    'EProp',
+    'OSTLRecurrent',
+    'OSTLFeedforward',
+    'OTPE',
+    'OTTT',
+    'OSTTP',
+    'FixedRandomFeedback',
+    'KappaFilter',
+    'PresynapticTrace',
+
+    # errors
+    'NotSupportedError',
+    'CompilationError',
+
+    # submodules
+    'nn',
+]
+
+
+# --- v0.1.x legacy shims: deprecated, served lazily with an access-time warning.
+# Each maps the public name -> migration replacement text. The shim classes still
+# work; new code should use the primitive-based ETP user-API instead.
+_DEPRECATED_LEGACY = {
+    'MatMulOp': 'braintrace.matmul (with a brainstate.ParamState)',
+    'ElemWiseOp': 'braintrace.element_wise',
+    'ConvOp': 'braintrace.conv',
+    'SpMatMulOp': 'braintrace.sparse_matmul',
+    'LoraOp': 'braintrace.lora_matmul',
+    'ETraceOp': 'the braintrace ETP primitive functions (matmul, conv, ...)',
+    'ETraceParam': 'brainstate.ParamState together with an ETP primitive function',
+    'ElemWiseParam': 'brainstate.ParamState together with braintrace.element_wise',
+    'NonTempParam': 'brainstate.ParamState with plain JAX ops (keeps the weight out of the ETP graph)',
+    'FakeETraceParam': 'a plain object with plain JAX ops',
+    'FakeElemWiseParam': 'a plain object with plain JAX ops',
+}
+
+
+def __getattr__(name):
+    if name in _DEPRECATED_LEGACY:
+        import warnings
+        warnings.warn(
+            f'braintrace.{name} is deprecated since 0.2.0 and will be removed in a '
+            f'future release; use {_DEPRECATED_LEGACY[name]} instead.',
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        from . import _legacy
+        return getattr(_legacy, name)
+    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')
+
+
+def __dir__():
+    return sorted(list(__all__) + list(_DEPRECATED_LEGACY))

{braintrace-0.1.2 → braintrace-0.2.1}/braintrace/_compatible_imports.py

@@ -29,11 +29,7 @@ __all__ = [
     'is_cond_primitive',
 ]
 
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
-
-else:
-    from jax.extend.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
+from brainstate._compatible_import import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
 
 
 def new_var(suffix, aval):

braintrace-0.2.1/braintrace/_compile.py

@@ -0,0 +1,153 @@
+# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Type, Union
+
+from ._etrace_algorithms import (
+    ETraceAlgorithm,
+    D_RTRL,
+    pp_prop,
+    EProp,
+    OSTLRecurrent,
+    OSTLFeedforward,
+    OTPE,
+    OTTT,
+    OSTTP,
+)
+
+__all__ = ['compile']
+
+# Canonical lowercase name (+ aliases) -> algorithm class. No bare ``ostl``
+# alias: the ambiguous OSTL factory was removed in 0.2.0, so callers pick
+# ``ostl_recurrent`` vs ``ostl_feedforward`` explicitly.
+_ALGORITHM_REGISTRY: dict[str, type[ETraceAlgorithm]] = {
+    'd_rtrl': D_RTRL,
+    'pp_prop': pp_prop,
+    'es_d_rtrl': pp_prop,
+    'esd_rtrl': pp_prop,
+    'eprop': EProp,
+    'e_prop': EProp,
+    'ostl_recurrent': OSTLRecurrent,
+    'ostl_feedforward': OSTLFeedforward,
+    'otpe': OTPE,
+    'ottt': OTTT,
+    'osttp': OSTTP,
+}
+
+
+def _resolve_algorithm(
+    algorithm: Union[str, Type[ETraceAlgorithm]]
+) -> Type[ETraceAlgorithm]:
+    """Resolve ``algorithm`` to an :class:`ETraceAlgorithm` subclass.
+
+    Parameters
+    ----------
+    algorithm : type or str
+        Either an :class:`ETraceAlgorithm` subclass (returned unchanged) or a
+        registered string name (case-insensitive), e.g. ``'D_RTRL'``,
+        ``'eprop'``, ``'ottt'``.
+
+    Returns
+    -------
+    type
+        The resolved :class:`ETraceAlgorithm` subclass.
+
+    Raises
+    ------
+    ValueError
+        If ``algorithm`` is a string that is not a registered name.
+    TypeError
+        If ``algorithm`` is a class that is not an ``ETraceAlgorithm`` subclass,
+        or is neither a class nor a string.
+    """
+    if isinstance(algorithm, type):
+        if issubclass(algorithm, ETraceAlgorithm):
+            return algorithm
+        raise TypeError(
+            f'algorithm class must be a subclass of ETraceAlgorithm, got {algorithm!r}.'
+        )
+    if isinstance(algorithm, str):
+        key = algorithm.strip().lower()
+        try:
+            return _ALGORITHM_REGISTRY[key]
+        except KeyError:
+            valid = ', '.join(sorted(_ALGORITHM_REGISTRY))
+            raise ValueError(
+                f'Unknown algorithm name {algorithm!r}. Valid names: {valid}. '
+                f'Or pass an ETraceAlgorithm subclass directly.'
+            )
+    raise TypeError(
+        f'algorithm must be an ETraceAlgorithm subclass or a registered string name, '
+        f'got {type(algorithm)}.'
+    )
+
+
+def compile(model, algorithm, *example_inputs, **options):
+    """Construct an online-learning algorithm for ``model`` and eagerly build its
+    eligibility-trace graph, returning a ready-to-``update`` learner.
+
+    Parameters
+    ----------
+    model : brainstate.nn.Module
+        The recurrent model. Its states must already be initialized, e.g. via
+        ``brainstate.nn.init_all_states(model)``.
+    algorithm : type or str
+        An :class:`ETraceAlgorithm` subclass, or a registered string name
+        (case-insensitive), e.g. ``'D_RTRL'``, ``'eprop'``, ``'ottt'``.
+    *example_inputs
+        Example call inputs (arrays / :class:`SingleStepData` /
+        :class:`MultiStepData`), matching what ``learner.update(...)`` will later
+        receive. Forwarded to :meth:`ETraceAlgorithm.compile_graph` to trace the
+        jaxpr graph. At least one is required.
+    **options
+        Keyword options forwarded to the algorithm constructor, e.g.
+        ``vjp_method``, ``leak``, ``fast_solve``, ``trace_dtype``, ``feedback``.
+
+    Returns
+    -------
+    ETraceAlgorithm
+        The compiled learner; call ``.update(*inputs)`` to train.
+
+    Raises
+    ------
+    ValueError
+        If ``algorithm`` is an unknown string name, or no ``example_inputs`` are
+        given.
+    TypeError
+        If ``algorithm`` is neither an ``ETraceAlgorithm`` subclass nor a string.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import braintrace
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>> model = MyRNN()
+        >>> brainstate.nn.init_all_states(model, batch_size=1)
+        >>> x0 = jnp.ones((3,))
+        >>> learner = braintrace.compile(model, 'D_RTRL', x0, vjp_method='multi-step')
+        >>> y = learner.update(x0)
+    """
+    cls = _resolve_algorithm(algorithm)
+    if len(example_inputs) == 0:
+        raise ValueError(
+            'compile() needs at least one example input to build the graph '
+            'eagerly, e.g. compile(model, "D_RTRL", x0). Pass the same inputs '
+            'you will give to learner.update(...).'
+        )
+    learner = cls(model, **options)
+    learner.compile_graph(*example_inputs)
+    return learner