{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0275870d",
   "metadata": {},
   "source": [
    "# e-prop: local eligibility and learning signals\n",
    "\n",
    "e-prop separates a synapse-local eligibility trace from a learning signal\n",
    "broadcast by the readout. The separation is the scientific content of the\n",
    "method; a generic loss curve alone would not demonstrate it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "465bd6db",
   "metadata": {},
   "source": [
    "## 1. Principle\n",
    "\n",
    "For synapse $i\\to j$, e-prop forms\n",
    "\n",
    "$$\\frac{d\\mathcal{L}}{dW_{ji}} \\approx\n",
    "  \\sum_t L_j^t\\,\\bar e_{ji}^t,\\qquad\n",
    "  \\bar e_{ji}^t=\\kappa\\bar e_{ji}^{t-1}+e_{ji}^t.$$\n",
    "\n",
    "$e_{ji}^t$ is local to the neuron/synapse dynamics. $L_j^t$ carries task\n",
    "information from the readout. Symmetric feedback uses the actual reverse-AD\n",
    "learning-signal direction. Random feedback uses a fixed projection and removes\n",
    "weight transport only partially in the current hook contract; it must not be\n",
    "described as numerically equivalent to symmetric feedback.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30ae2edd",
   "metadata": {},
   "source": [
    "## 2. Applicability and exclusions\n",
    "\n",
    "**Use e-prop for** recurrent LIF/ALIF-style networks when local traces and an\n",
    "explicit learning-signal approximation are the intended inductive bias.\n",
    "\n",
    "**Do not infer** exact BPTT gradients for arbitrary recurrent coupling. The\n",
    "$\\kappa$ filter changes temporal credit assignment, while random feedback adds\n",
    "a separate approximation and requires an explicit key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "891d66ca",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:30:57.921306Z",
     "iopub.status.busy": "2026-08-07T14:30:57.921306Z",
     "iopub.status.idle": "2026-08-07T14:31:40.783239Z",
     "shell.execute_reply": "2026-08-07T14:31:40.781552Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainstate\n",
    "import braintools\n",
    "import braintrace\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "brainstate.random.seed(21)\n",
    "\n",
    "\n",
    "class SmallEPropSNN(brainstate.nn.Module):\n",
    "    def __init__(self, n_in=2, n_rec=6):\n",
    "        super().__init__()\n",
    "        self.n_rec = n_rec\n",
    "        self.w_in = brainstate.ParamState(\n",
    "            0.3 * brainstate.random.randn(n_in, n_rec)\n",
    "        )\n",
    "        self.w_rec = brainstate.ParamState(\n",
    "            0.1 * brainstate.random.randn(n_rec, n_rec)\n",
    "        )\n",
    "        self.w_out = brainstate.ParamState(\n",
    "            0.2 * brainstate.random.randn(n_rec, 1)\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.z = 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.z.value = jnp.zeros_like(self.z.value)\n",
    "\n",
    "    def update(self, x):\n",
    "        drive = (\n",
    "            braintrace.matmul(x, self.w_in.value)\n",
    "            + braintrace.matmul(self.z.value, self.w_rec.value)\n",
    "        )\n",
    "        voltage = 0.9 * self.v.value + drive - self.z.value\n",
    "        spike = self.surrogate(voltage - 0.45)\n",
    "        self.v.value = voltage\n",
    "        self.z.value = spike\n",
    "        return spike @ self.w_out.value\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d7b683",
   "metadata": {},
   "source": [
    "## 3. Compile and expose the filter coordinate\n",
    "\n",
    "The sequence driver advances the SNN once per call. This example uses symmetric\n",
    "feedback so the effect under inspection is the $\\kappa$ eligibility filter,\n",
    "not a simultaneous random-feedback approximation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a48cb265",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:31:40.785645Z",
     "iopub.status.busy": "2026-08-07T14:31:40.785645Z",
     "iopub.status.idle": "2026-08-07T14:31:43.822249Z",
     "shell.execute_reply": "2026-08-07T14:31:43.821622Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trace filter: kappa\n",
      "kappa: 0.8\n",
      "learning signal: symmetric\n",
      "masked loss: 0.0586\n",
      "finite nonzero trace: True\n",
      "finite nonzero gradient: True\n"
     ]
    }
   ],
   "source": [
    "inputs = brainstate.random.bernoulli(\n",
    "    0.35, size=(16, 2)\n",
    ").astype(jnp.float32)\n",
    "targets = jnp.full((16, 1), 0.25)\n",
    "mask = (jnp.arange(16) >= 4).astype(jnp.float32)\n",
    "\n",
    "model = SmallEPropSNN()\n",
    "brainstate.nn.init_all_states(model)\n",
    "learner = braintrace.compile(\n",
    "    model,\n",
    "    braintrace.EProp,\n",
    "    inputs[0],\n",
    "    feedback=\"symmetric\",\n",
    "    kappa_filter_decay=0.8,\n",
    ")\n",
    "\n",
    "\n",
    "def loss_step(x, target):\n",
    "    return jnp.mean((learner(x) - target) ** 2)\n",
    "\n",
    "\n",
    "brainstate.nn.reset_all_states(model)\n",
    "learner.reset_state()\n",
    "grads, loss = learner.etrace_grad(\n",
    "    inputs,\n",
    "    targets,\n",
    "    step_fn=loss_step,\n",
    "    mask=mask,\n",
    "    loss_output=\"scalar\",\n",
    "    return_value=True,\n",
    ")\n",
    "gradient_norm = jnp.sqrt(sum(jnp.sum(g * g) for g in grads.values()))\n",
    "etraces = learner.get_etrace_of(model.w_rec)\n",
    "trace_norm = jnp.sqrt(sum(\n",
    "    jnp.sum(leaf * leaf) for leaf in jax.tree.leaves(etraces)\n",
    "))\n",
    "\n",
    "print(\"trace filter:\", learner.config.trace_filter)\n",
    "print(\"kappa:\", learner.config.kappa)\n",
    "print(\"learning signal:\", learner.config.learning_signal)\n",
    "print(f\"masked loss: {float(loss):.4f}\")\n",
    "print(\"finite nonzero trace:\", bool(jnp.isfinite(trace_norm) & (trace_norm > 0)))\n",
    "print(\"finite nonzero gradient:\", bool(jnp.isfinite(gradient_norm) & (gradient_norm > 0)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d3642d4",
   "metadata": {},
   "source": [
    "## 4. Interpretation and limits\n",
    "\n",
    "The printed configuration verifies which learning-rule coordinates were\n",
    "compiled. The nonzero recurrent-weight trace shows that the local eligibility\n",
    "state was advanced, while the finite gradient checks its contraction through\n",
    "the public sequence path. Neither check isolates the numerical effect of\n",
    "$\\kappa$. A causal\n",
    "comparison would hold model, sequence, feedback, and parameters fixed while\n",
    "varying only $\\kappa$, then use a finite-window oracle appropriate to the\n",
    "learning-rule claim.\n",
    "\n",
    "## Reference and API\n",
    "\n",
    "- Bellec et al., \"A solution to the learning dilemma for recurrent networks of\n",
    "  spiking neurons,\" *Nature Communications* 11, 3625 (2020),\n",
    "  [doi:10.1038/s41467-020-17236-y](https://doi.org/10.1038/s41467-020-17236-y).\n",
    "- API: {class}`braintrace.EProp`, {func}`braintrace.compile`,\n",
    "  {meth}`braintrace.SequenceDriverMixin.etrace_grad`.\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
}