{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3cdb6f4b",
   "metadata": {},
   "source": [
    "# pp-prop: input/output-factorized online traces\n",
    "\n",
    "pp-prop is the public name of BrainTrace's input/output-factorized estimator\n",
    "(`ES_D_RTRL` remains a compatibility alias). This example uses a recurrent\n",
    "spiking model because a matched MiniGRU loss curve would not demonstrate the\n",
    "method's intended linear-memory SNN regime.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de939794",
   "metadata": {},
   "source": [
    "## 1. Principle and approximation\n",
    "\n",
    "Instead of storing a parameter-shaped eligibility tensor, pp-prop maintains an\n",
    "input-side factor $x_t$ and an output/hidden-side factor $f_t$. Their smoothed\n",
    "outer-product contraction estimates the parameter trace:\n",
    "\n",
    "$$\\widehat{G}_t \\approx \\bar{x}_t \\otimes \\bar{f}_t,\\qquad\n",
    "\\bar{x}_t=\\rho\\bar{x}_{t-1}+(1-\\rho)x_t.$$\n",
    "\n",
    "The analogous recursion is applied to the output-side factor. This changes the\n",
    "memory dependence, but discards correlations that cannot be represented by the\n",
    "factorization. The decay is part of the estimator, not merely an optimizer\n",
    "hyperparameter.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d99dbd",
   "metadata": {},
   "source": [
    "## 2. Applicability\n",
    "\n",
    "**Use pp-prop when** a recurrent SNN needs an input/output-factorized trace and\n",
    "the full parameter-dimensional trace is too costly.\n",
    "\n",
    "**Avoid unqualified claims when** strong cross-neuron correlations dominate,\n",
    "the decay has not been validated, or exact BPTT gradients are required. Windowed\n",
    "`etrace_grad` is also inappropriate for the current IO-factorized bias\n",
    "correction; drive one time step at a time as below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a7556cf9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:13:16.057646Z",
     "iopub.status.busy": "2026-08-07T14:13:16.057646Z",
     "iopub.status.idle": "2026-08-07T14:13:19.621254Z",
     "shell.execute_reply": "2026-08-07T14:13:19.620322Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainstate\n",
    "import braintools\n",
    "import braintrace\n",
    "import jax.numpy as jnp\n",
    "\n",
    "brainstate.random.seed(11)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "432e6b75",
   "metadata": {},
   "source": [
    "## 3. Minimal recurrent spiking model\n",
    "\n",
    "Each call computes one membrane update and one surrogate spike. The spike is\n",
    "then used once by the recurrent path and once by the readout. This one-call\n",
    "contract matters: calling a neuron twice would advance its state twice while\n",
    "`etrace_grad` believes only one logical time step elapsed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6a4fc635",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:13:19.625011Z",
     "iopub.status.busy": "2026-08-07T14:13:19.624258Z",
     "iopub.status.idle": "2026-08-07T14:13:19.636904Z",
     "shell.execute_reply": "2026-08-07T14:13:19.634859Z"
    }
   },
   "outputs": [],
   "source": [
    "class RecurrentLIF(brainstate.nn.Module):\n",
    "    def __init__(self, n_in=3, n_rec=8, n_out=1):\n",
    "        super().__init__()\n",
    "        self.n_rec = n_rec\n",
    "        self.w_in = brainstate.ParamState(\n",
    "            0.35 * brainstate.random.randn(n_in, n_rec)\n",
    "        )\n",
    "        self.w_rec = brainstate.ParamState(\n",
    "            0.12 * brainstate.random.randn(n_rec, n_rec)\n",
    "        )\n",
    "        self.w_out = brainstate.ParamState(\n",
    "            0.2 * brainstate.random.randn(n_rec, n_out)\n",
    "        )\n",
    "        self.surrogate = braintools.surrogate.ReluGrad()\n",
    "\n",
    "    def init_state(self, **kwargs):\n",
    "        self.v = brainstate.HiddenState(jnp.zeros(self.n_rec))\n",
    "        self.spike = brainstate.HiddenState(jnp.zeros(self.n_rec))\n",
    "\n",
    "    def reset_state(self, **kwargs):\n",
    "        self.v.value = jnp.zeros_like(self.v.value)\n",
    "        self.spike.value = jnp.zeros_like(self.spike.value)\n",
    "\n",
    "    def update(self, x):\n",
    "        current = (\n",
    "            braintrace.matmul(x, self.w_in.value)\n",
    "            + braintrace.matmul(self.spike.value, self.w_rec.value)\n",
    "        )\n",
    "        voltage = 0.85 * self.v.value + current - self.spike.value\n",
    "        spike = self.surrogate(voltage - 0.5)\n",
    "        self.v.value = voltage\n",
    "        self.spike.value = spike\n",
    "        return voltage @ self.w_out.value\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d05fb5",
   "metadata": {},
   "source": [
    "## 4. Compile one time step and align the objective\n",
    "\n",
    "The compilation sample is `inputs[0]`, one feature vector, not the whole\n",
    "sequence. A warm-up mask gates the loss while still advancing both neuron and\n",
    "eligibility states. Evaluation uses the same mask and per-step squared error as\n",
    "training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e5bcc8f4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:13:19.640904Z",
     "iopub.status.busy": "2026-08-07T14:13:19.639902Z",
     "iopub.status.idle": "2026-08-07T14:13:23.638547Z",
     "shell.execute_reply": "2026-08-07T14:13:23.637543Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "initial masked loss: 0.1687\n",
      "final masked loss:   0.0467\n",
      "finite trajectory: True\n",
      "loss decreased: True\n",
      "compiled relations: 2\n"
     ]
    }
   ],
   "source": [
    "n_steps = 24\n",
    "inputs = brainstate.random.bernoulli(\n",
    "    0.3, size=(n_steps, 3)\n",
    ").astype(jnp.float32)\n",
    "targets = jnp.full((n_steps, 1), 0.4)\n",
    "loss_mask = (jnp.arange(n_steps) >= 6).astype(jnp.float32)\n",
    "\n",
    "model = RecurrentLIF()\n",
    "brainstate.nn.init_all_states(model)\n",
    "learner = braintrace.compile(\n",
    "    model,\n",
    "    braintrace.pp_prop,\n",
    "    inputs[0],\n",
    "    decay_or_rank=0.5,\n",
    ")\n",
    "optimizer = braintools.optim.Adam(2e-3)\n",
    "optimizer.register_trainable_weights(learner.param_states)\n",
    "\n",
    "\n",
    "def reset_sequence():\n",
    "    brainstate.nn.reset_all_states(model)\n",
    "    learner.reset_state()\n",
    "\n",
    "\n",
    "def step_loss(x, target):\n",
    "    prediction = learner(x)\n",
    "    return jnp.mean((prediction - target) ** 2)\n",
    "\n",
    "\n",
    "def masked_objective():\n",
    "    reset_sequence()\n",
    "    outputs = learner.etrace_evolve(inputs, return_outputs=True)\n",
    "    losses = jnp.mean((outputs - targets) ** 2, axis=1)\n",
    "    return jnp.sum(losses * loss_mask) / jnp.sum(loss_mask)\n",
    "\n",
    "\n",
    "def train_epoch(_):\n",
    "    reset_sequence()\n",
    "    grads, objective = learner.etrace_grad(\n",
    "        inputs,\n",
    "        targets,\n",
    "        step_fn=step_loss,\n",
    "        mask=loss_mask,\n",
    "        reduction=\"mean\",\n",
    "        loss_output=\"scalar\",\n",
    "        return_value=True,\n",
    "    )\n",
    "    optimizer.update(brainstate.nn.clip_grad_norm(grads, 1.0))\n",
    "    return objective\n",
    "\n",
    "\n",
    "initial_loss = masked_objective()\n",
    "training_losses = brainstate.transform.for_loop(\n",
    "    train_epoch, jnp.arange(30)\n",
    ")\n",
    "final_loss = masked_objective()\n",
    "\n",
    "print(f\"initial masked loss: {float(initial_loss):.4f}\")\n",
    "print(f\"final masked loss:   {float(final_loss):.4f}\")\n",
    "print(\"finite trajectory:\", bool(jnp.all(jnp.isfinite(training_losses))))\n",
    "print(\"loss decreased:\", bool(final_loss < initial_loss))\n",
    "print(\"compiled relations:\", len(learner.graph.hidden_param_op_relations))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63693957",
   "metadata": {},
   "source": [
    "## 5. What the check establishes\n",
    "\n",
    "The run verifies the public compile/sequence API, single-step state evolution,\n",
    "finite factorized updates, a shared train/evaluation window, and descent on one\n",
    "fixed task. It does **not** measure factorization error against BPTT. That needs\n",
    "a reduced finite-window gradient oracle for the intended model and decay.\n",
    "\n",
    "## References and API\n",
    "\n",
    "- Wang et al., \"Model-agnostic linear-memory online learning in spiking neural\n",
    "  networks,\" *Nature Communications* (2026),\n",
    "  [doi:10.1038/s41467-026-68453-w](https://doi.org/10.1038/s41467-026-68453-w).\n",
    "- API: {class}`braintrace.pp_prop`, {class}`braintrace.IODimVjpAlgorithm`,\n",
    "  {func}`braintrace.compile`.\n",
    "- Continue with {doc}`snn_online_learning` for the batched workflow and\n",
    "  {doc}`eprop` for local eligibility traces with an explicit learning signal.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}