Correct D-RTRL/pp_prop docs and decay_or_rank semantics

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

braintrace/__init___test.py

@@ -171,6 +171,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):

braintrace/_algorithm/io_dim_vjp.py

@@ -639,10 +639,11 @@ class IODimVjpAlgorithm(ETraceVjpAlgorithm):
         The model function, which receives the input arguments and returns the
         model output.
     decay_or_rank : float or int
-        The exponential smoothing factor for the eligibility trace. If a float,
-        it is the decay factor and should be in the range :math:`(0, 1)`. If an
-        integer, it is the number of approximation ranks for the algorithm and
-        should be greater than 0.
+        Parameterization of the exponential smoothing factor. A float in
+        :math:`(0, 1)` is used directly as the decay. A positive integer
+        :math:`r` is converted to :math:`\alpha=(r-1)/(r+1)`. The integer form
+        does not allocate :math:`r` separate trace factors; both forms use the
+        same input/output-factorized trace structure.
     vjp_method : str, optional
         The method for computing the VJP. It should be either ``"single-step"``
         or ``"multi-step"``.
@@ -695,8 +696,8 @@ class IODimVjpAlgorithm(ETraceVjpAlgorithm):
     the outer product of two exponentially-smoothed *vectors* — one over the
     input dimension and one over the output dimension. Storing the two factors
     instead of the full matrix drops the memory from :math:`O(I\cdot O)` to
-    :math:`O(I+O)` per layer. The decay :math:`\alpha` (equivalently an
-    approximation rank) controls how much temporal history the factored trace
+    :math:`O(I+O)` per layer. The decay :math:`\alpha` (or its integer
+    parameterization) controls how much temporal history the factored trace
     retains; the bias of the exponential estimator is corrected at solve time.
 
     This algorithm has :math:`O(BI+BO)` memory complexity and :math:`O(BIO)`

braintrace/nn/__init__.py

@@ -38,7 +38,7 @@
 Activation, normalisation and pooling layers are intentionally not
 re-implemented here; accessing them through ``braintrace.nn`` emits a
 :class:`DeprecationWarning` and forwards to :mod:`brainstate.nn` /
-:mod:`brainstate.state`.
+:mod:`brainpy.state`.
 """
 
 from __future__ import annotations
@@ -70,12 +70,12 @@ def __getattr__(name: str) -> Any:
     import warnings
     if name in ['IF', 'LIF', 'ALIF', 'Expon', 'Alpha', 'DualExpon', 'STP', 'STD']:
         warnings.warn(
-            f'braintrace.nn.{name} is deprecated. Use brainstate.state.{name} instead.',
+            f'braintrace.nn.{name} is deprecated. Use brainpy.state.{name} instead.',
             DeprecationWarning,
             stacklevel=2
         )
-        import brainstate.state
-        return getattr(brainstate.state, name)
+        import brainpy.state
+        return getattr(brainpy.state, name)
 
     if name in [
         'ReLU', 'RReLU', 'Hardtanh', 'ReLU6', 'Sigmoid', 'Hardsigmoid',

docs/apis/algorithms.rst

@@ -101,12 +101,14 @@ is the state these algorithms carry across time.
    EligibilityTrace
 
 
-D-RTRL — Parameter Dimension (exact)
-------------------------------------
+D-RTRL — Parameter-dimensional estimator
+-----------------------------------------
 
-Decoupled Real-Time Recurrent Learning with a diagonal approximation of the
-hidden-to-hidden Jacobian. Memory complexity :math:`O(B \cdot |\theta|)`, where
-:math:`B` is the batch size and :math:`|\theta|` the number of parameters.
+Diagonal Real-Time Recurrent Learning uses a diagonal approximation of the
+hidden-to-hidden Jacobian. Memory complexity is
+:math:`O(B \cdot |\theta|)`, where :math:`B` is the batch size and
+:math:`|\theta|` the number of parameters. It is not generally
+gradient-equivalent to BPTT outside the assumptions of that approximation.
 
 .. math::
 
@@ -131,13 +133,14 @@ hidden-to-hidden Jacobian. Memory complexity :math:`O(B \cdot |\theta|)`, where
 :class:`ParamDimVjpAlgorithm`.
 
 
-ES-D-RTRL — Input/Output Dimension (exact)
-------------------------------------------
+pp-prop — Input/output-factorized estimator
+-------------------------------------------
 
-The Event-Synchronized D-RTRL algorithm factorizes the eligibility trace into
-input and output components with exponential smoothing, reducing memory to
-:math:`O(B(I + O))`, where :math:`I` and :math:`O` are the input and output
-dimensions.
+``pp_prop`` (historically exposed as ``ES_D_RTRL``) factorizes the eligibility
+trace into input and output components with exponential smoothing, reducing
+memory to :math:`O(B(I + O))`, where :math:`I` and :math:`O` are the input and
+output dimensions. An integer ``decay_or_rank`` value parameterizes the decay;
+it does not allocate multiple rank factors.
 
 .. math::
 

docs/apis/nn.rst

@@ -18,7 +18,7 @@ for gate operations.
    Activation, normalization, and pooling layers are intentionally **not**
    re-implemented here. Accessing them through ``braintrace.nn`` (e.g.
    ``braintrace.nn.LayerNorm``) emits a :class:`DeprecationWarning` and forwards
-   to ``brainstate.nn`` / ``brainstate.state``; use those packages directly.
+   to ``brainstate.nn`` / ``brainpy.state``; use those packages directly.
 
 
 Linear Layers

docs/quickstart/rnn_online_learning.ipynb

@@ -13,7 +13,7 @@
     "\n",
     "1. Define the **copying task**, a standard benchmark for testing sequential memory in RNNs.\n",
     "2. Build a GRU model using `braintrace.nn` components.\n",
-    "3. Train the model with **D-RTRL** (Decoupled Real-Time Recurrent Learning), an online learning algorithm that computes approximate gradients without storing the full computation graph.\n",
+    "3. Train the model with **D-RTRL** (Diagonal Real-Time Recurrent Learning), an online learning algorithm that computes approximate gradients without storing the full computation graph.\n",
     "4. Compare the online learning approach with standard **Backpropagation Through Time (BPTT)**.\n",
     "\n",
     "Online learning is especially useful when:\n",
@@ -180,7 +180,7 @@
    "cell_type": "markdown",
    "id": "a1b2c3d4e5f60008",
    "metadata": {},
-   "source": "## 4. Online Training with D-RTRL\n\nD-RTRL (Decoupled Real-Time Recurrent Learning) is an online learning algorithm provided by `braintrace`. Unlike BPTT, which requires storing the entire computation graph across all time steps, D-RTRL computes gradients incrementally at each time step using **eligibility traces**.\n\nThe key steps in the online training loop are:\n\n1. **Compile the model**: `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initialises hidden states, compiles the eligibility-trace graph, and returns a vmapped learner — all in one call.\n2. **Warm-up phase**: Run the model forward (without learning) to let the hidden states and eligibility traces stabilize.\n3. **Learning phase**: At each time step, compute the gradient of the current loss with respect to the parameters, and accumulate gradients over time.\n4. **Parameter update**: After processing the full sequence, apply the accumulated gradients to update the parameters.\n\nThe `D_RTRL` class wraps the model and handles the eligibility trace bookkeeping automatically."
+   "source": "## 4. Online Training with D-RTRL\n\nD-RTRL (Diagonal Real-Time Recurrent Learning) is an online learning algorithm provided by `braintrace`. Unlike BPTT, which requires storing the entire computation graph across all time steps, D-RTRL computes approximate gradients incrementally at each time step using **eligibility traces**. It is not generally gradient-equivalent to BPTT outside the assumptions of its diagonal Jacobian approximation.\n\nThe key steps in the online training loop are:\n\n1. **Compile the model**: `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initialises hidden states, compiles the eligibility-trace graph, and returns a vmapped learner — all in one call.\n2. **Warm-up phase**: Run the model forward (without learning) to let the hidden states and eligibility traces stabilize.\n3. **Learning phase**: At each time step, compute the gradient of the current loss with respect to the parameters, and accumulate gradients over time.\n4. **Parameter update**: After processing the full sequence, apply the accumulated gradients to update the parameters.\n\nThe `D_RTRL` class wraps the model and handles the eligibility trace bookkeeping automatically."
   },
   {
    "cell_type": "code",
@@ -467,4 +467,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

docs/quickstart/snn_online_learning.ipynb

@@ -21,7 +21,7 @@
     "\n",
     "**Online learning** is a natural fit for SNNs because they process inputs sequentially, one time step at a time. Instead of storing the entire computation graph for backpropagation through time (BPTT), online learning algorithms update weight gradients incrementally at each time step. This eliminates the need to unroll the network over the full sequence length, resulting in constant memory usage with respect to sequence length.\n",
     "\n",
-    "In this tutorial, we will use **ES-D-RTRL** (Eligibility-trace Scalable Decoupled Real-Time Recurrent Learning), an efficient online learning algorithm provided by `braintrace`. ES-D-RTRL factorizes the eligibility trace into input and output components, achieving **O(B(I+O))** memory complexity (where B is batch size, I is input dimension, and O is output dimension). This makes it highly scalable for large spiking networks.\n",
+    "In this tutorial, we use `braintrace.pp_prop` (historically exposed as `braintrace.ES_D_RTRL`), an online estimator that factorizes the eligibility trace into input and output components. Its trace memory is **O(B(I+O))** (where B is batch size, I is input dimension, and O is output dimension), subject to the documented model and operator assumptions.\n",
     "\n",
     "**What you will learn:**\n",
     "- How to build an SNN model using `brainstate` neurons and `braintrace.nn` layers\n",
@@ -198,7 +198,7 @@
    "cell_type": "markdown",
    "id": "43f58871",
    "metadata": {},
-   "source": "## 3. Training with ES-D-RTRL\n\nWe set up online learning using `braintrace.compile` with `braintrace.pp_prop` (also exposed as `braintrace.ES_D_RTRL` and `braintrace.IODimVjpAlgorithm`). The `decay_or_rank` parameter controls the trace approximation:\n\n* **`decay_or_rank=float` in (0, 1]** -- exponentially-smoothed trace. The value is the decay factor applied per step; `0.99` is a common choice. Memory cost: `O(B * (I + O))` per layer.\n* **`decay_or_rank=int >= 1`** -- low-rank trace. The integer is the rank used to factorise the trace. Memory cost: `O(B * rank * (I + O))`.\n\nPick the **decay form** when you want a single hyper-parameter that you can sweep cheaply, and the **rank form** when you want a tunable accuracy/memory trade-off independent of any time-scale assumption.\n\n`braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.99)` initialises per-sample states, builds the ETP graph, and returns a vmapped learner — no separate `init_all_states`, `compile_graph`, or `Vmap` calls are needed.\n\n`braintrace.D_RTRL` is the alternative algorithm; it stores the *full* parameter-dimension trace (`O(B * theta)`) and is exact rather than approximate. Use `D_RTRL` when memory permits and `ES_D_RTRL` for larger networks."
+   "source": "## 3. Training with ES-D-RTRL\n\nWe set up online learning using `braintrace.compile` with `braintrace.pp_prop` (also exposed as `braintrace.ES_D_RTRL` and `braintrace.IODimVjpAlgorithm`). The `decay_or_rank` parameter controls the trace approximation:\n\n* **`decay_or_rank=float` in (0, 1)** -- exponentially-smoothed trace. The value is the decay factor applied per step; `0.99` is a common choice. Memory cost: `O(B * (I + O))` per layer.\n* **`decay_or_rank=int >= 1`** -- an alternative parameterization of the same decay, using `decay = (rank - 1) / (rank + 1)`. It does not allocate multiple rank factors, so memory remains `O(B * (I + O))`.\n\nUse the float form to set the decay directly, or the integer form to select the corresponding decay through the documented conversion. Neither form creates an independent rank-versus-memory trade-off.\n\n`braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.99)` initialises per-sample states, builds the ETP graph, and returns a vmapped learner — no separate `init_all_states`, `compile_graph`, or `Vmap` calls are needed.\n\n`braintrace.D_RTRL` is the alternative parameter-dimensional estimator. It stores parameter-shaped eligibility traces and can use substantially more memory; neither estimator is generally gradient-equivalent to BPTT outside its documented assumptions."
   },
   {
    "cell_type": "code",
@@ -331,7 +331,7 @@
    "cell_type": "markdown",
    "id": "cdbdc871",
    "metadata": {},
-   "source": "## 4. Key Differences: D-RTRL vs ES-D-RTRL\n\nBrainTrace provides two main online learning algorithms. Understanding their trade-offs helps you choose the right one for your application.\n\n| Aspect | D-RTRL (`ParamDimVjpAlgorithm`) | ES-D-RTRL (`ES_D_RTRL`) |\n|---|---|---|\n| **Eligibility Trace** | Full trace per weight parameter | Factorized into input/output components |\n| **Memory Complexity** | O(B * theta) where theta = total parameters | O(B * (I + O)) where I = input dim, O = output dim |\n| **Computation** | Exact gradient computation | Approximation via low-rank factorization |\n| **Scalability** | Suitable for small networks | Scales to large networks (hundreds/thousands of neurons) |\n| **Use Cases** | Research requiring exact gradients | Practical SNN training, large-scale networks |\n| **BrainTrace API** | `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B)` | `braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, decay_or_rank=0.99)` |\n\n### When to use which?\n\n- **D-RTRL** stores the full eligibility trace for each weight, giving exact online gradients. However, this requires O(B * theta) memory, which grows linearly with the number of parameters. For a network with N hidden neurons and a recurrent weight matrix of size N x N, the trace has N^4 entries per sample. This limits D-RTRL to small networks (typically < 100 neurons).\n\n- **ES-D-RTRL** factorizes the eligibility trace into input and output components, reducing memory to O(B * (I + O)). The `decay_or_rank` parameter controls the approximation: a float value (e.g., 0.99) sets the trace decay factor, while an integer value sets the rank of the low-rank approximation. ES-D-RTRL is the recommended choice for SNNs, where networks often have hundreds or thousands of neurons.\n\nBoth algorithms are accessed via `braintrace.compile` — just swap the algorithm class and pass its hyperparameters as keyword arguments:\n\n```python\n# D-RTRL (exact, high memory)\nlearner = braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B)\n\n# ES-D-RTRL (approximate, scalable)\nlearner = braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, decay_or_rank=0.99)\n```"
+   "source": "## 4. Key differences: D-RTRL vs pp-prop\n\nBrainTrace provides two parameterizations of online eligibility-trace estimation. Their gradients and memory costs differ, and neither is generally equivalent to BPTT outside its documented assumptions.\n\n| Aspect | D-RTRL (`ParamDimVjpAlgorithm`) | pp-prop (`pp_prop`, historically `ES_D_RTRL`) |\n|---|---|---|\n| **Eligibility trace** | Parameter-shaped trace with a diagonal hidden-Jacobian approximation | Input/output-factorized trace with exponential smoothing |\n| **Trace memory** | O(B * theta), where theta is the number of traced parameters | O(B * (I + O)), where I and O are input and output dimensions |\n| **Main approximation** | Diagonal hidden-to-hidden Jacobian structure | Input/output factorization plus exponential smoothing |\n| **Recommended use** | When parameter-shaped trace memory is feasible | When lower trace memory is required and the factorization assumptions are acceptable |\n| **BrainTrace API** | `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B)` | `braintrace.compile(model, braintrace.pp_prop, x0, batch_size=B, decay_or_rank=0.99)` |\n\nA float `decay_or_rank` value sets the pp-prop decay directly. A positive integer selects the corresponding decay through `decay = (rank - 1) / (rank + 1)`; it does not allocate that many low-rank factors. Validate either estimator against BPTT on a reduced version of the intended model before relying on gradient fidelity.\n\n```python\n# D-RTRL: parameter-shaped trace\nlearner = braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B)\n\n# pp-prop: input/output-factorized trace\nlearner = braintrace.compile(model, braintrace.pp_prop, x0, batch_size=B, decay_or_rank=0.99)\n```"
   },
   {
    "cell_type": "markdown",
@@ -361,4 +361,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

docs/tutorials/drtrl.md

@@ -46,7 +46,7 @@ so both parameter sets end up in the trace.
   init + compile in a `vmap_new_states` scope, then `brainstate.nn.Vmap`.
   **Default choice.**
 - [`03-batching-batched.py`](../../examples/drtrl/03-batching-batched.py) —
-  `D_RTRL(..., mode=brainstate.mixin.Batching())`, init once with
+  `D_RTRL(...)`, init once with
   `batch_size=...`. Pick when the algorithm needs the batch axis exposed.
 
 ## 5. VJP methods

examples/003-snn-memory-and-speed-evaluation-all.py

@@ -355,9 +355,9 @@ def _compile_etrace_function(self, input_info):
         # kept manual: uses vmap_init_all_states with state_tag='new' and Batching() mode;
         # braintrace.compile uses init_all_states which is incompatible with this vmap state scheme
         if self.args.method == 'expsm_diag':
-            model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay, mode=brainstate.mixin.Batching())
+            model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay)
         elif self.args.method == 'diag':
-            model = braintrace.D_RTRL(self.target, mode=brainstate.mixin.Batching())
+            model = braintrace.D_RTRL(self.target)
         else:
             raise ValueError(f'Unknown online learning methods: {self.args.method}.')
 

examples/003-snn-memory-and-speed-evaluation-batched.py

@@ -479,9 +479,9 @@ def _compile_etrace_function(self, input_info):
         # kept manual: uses vmap_init_all_states with state_tag='new' and Batching() mode;
         # braintrace.compile uses init_all_states which is incompatible with this vmap state scheme
         if self.args.method == 'expsm_diag':
-            model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay, mode=brainstate.mixin.Batching())
+            model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay)
         elif self.args.method == 'diag':
-            model = braintrace.D_RTRL(self.target, mode=brainstate.mixin.Batching())
+            model = braintrace.D_RTRL(self.target)
         else:
             raise ValueError(f'Unknown online learning methods: {self.args.method}.')
 

examples/100-gru-on-copying-task.py

@@ -129,7 +129,7 @@ def batch_train(self, inputs, target):
         elif self.batch_train_method == 'batch':
             # kept manual: re-initializes states inside @jit every batch; braintrace.compile must live outside jit
             model = braintrace.ParamDimVjpAlgorithm(
-                self.target, vjp_method=self.vjp_method, mode=brainstate.mixin.Batching())
+                self.target, vjp_method=self.vjp_method)
             brainstate.nn.init_all_states(self.target, batch_size=inputs.shape[1])
             model.compile_graph(inputs[0])