{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a4d8370",
   "metadata": {},
   "source": [
    "# OSTL: with-H and without-H regimes\n",
    "\n",
    "Online Spatio-Temporal Learning is not one interchangeable preset. BrainTrace\n",
    "exposes the recurrent **with-H** and feedforward **without-H** regimes as\n",
    "separate classes because they retain different temporal terms.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6e57fd0",
   "metadata": {},
   "source": [
    "## 1. Principle\n",
    "\n",
    "The recurrent regime propagates\n",
    "\n",
    "$$\\epsilon_t = H_t\\epsilon_{t-1}+F_t,$$\n",
    "\n",
    "where $H_t=\\partial h_t/\\partial h_{t-1}$ is the hidden-to-hidden Jacobian and\n",
    "$F_t$ is the instantaneous parameter contribution. The without-H regime drops\n",
    "$H_t\\epsilon_{t-1}$ and keeps the spatial/instantaneous contribution. In\n",
    "BrainTrace, `OSTLRecurrent` uses a coupled per-parameter trace;\n",
    "`OSTLFeedforward` uses an IO-factorized trace with negligible temporal decay.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8495bab",
   "metadata": {},
   "source": [
    "## 2. Model-algorithm matching\n",
    "\n",
    "- **With-H:** use for a recurrent layer when the retained block structure\n",
    "  matches the model's hidden coupling.\n",
    "- **Without-H:** use for feedforward SNN dynamics where no recurrent Jacobian\n",
    "  should be propagated.\n",
    "- Do not apply the without-H rule to a recurrent model merely because it is\n",
    "  cheaper; that changes the estimator's mathematical target.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "411b72ae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:15:12.755821Z",
     "iopub.status.busy": "2026-08-07T14:15:12.755821Z",
     "iopub.status.idle": "2026-08-07T14:15:16.705441Z",
     "shell.execute_reply": "2026-08-07T14:15:16.703616Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainstate\n",
    "import braintrace\n",
    "import jax.numpy as jnp\n",
    "\n",
    "brainstate.random.seed(31)\n",
    "\n",
    "\n",
    "class OSTLModel(brainstate.nn.Module):\n",
    "    def __init__(self, recurrent):\n",
    "        super().__init__()\n",
    "        self.recurrent = recurrent\n",
    "        self.w_in = brainstate.ParamState(\n",
    "            0.25 * brainstate.random.randn(2, 5)\n",
    "        )\n",
    "        if recurrent:\n",
    "            self.w_rec = brainstate.ParamState(\n",
    "                0.1 * brainstate.random.randn(5, 5)\n",
    "            )\n",
    "        self.w_out = brainstate.ParamState(\n",
    "            0.2 * brainstate.random.randn(5, 1)\n",
    "        )\n",
    "\n",
    "    def init_state(self, **kwargs):\n",
    "        self.h = brainstate.HiddenState(jnp.zeros(5))\n",
    "\n",
    "    def reset_state(self, **kwargs):\n",
    "        self.h.value = jnp.zeros_like(self.h.value)\n",
    "\n",
    "    def update(self, x):\n",
    "        drive = braintrace.matmul(x, self.w_in.value)\n",
    "        if self.recurrent:\n",
    "            drive = drive + braintrace.matmul(self.h.value, self.w_rec.value)\n",
    "        self.h.value = jnp.tanh(drive)\n",
    "        return self.h.value @ self.w_out.value\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "792d29a1",
   "metadata": {},
   "source": [
    "## 3. Compile each regime on a matching model\n",
    "\n",
    "Both examples use the same public compile and sequence interfaces. They are not\n",
    "a head-to-head accuracy benchmark because the model structures intentionally\n",
    "differ.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8d3b79c5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:15:16.708038Z",
     "iopub.status.busy": "2026-08-07T14:15:16.708038Z",
     "iopub.status.idle": "2026-08-07T14:15:19.551778Z",
     "shell.execute_reply": "2026-08-07T14:15:19.550771Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "with-H recurrence scope: coupled\n",
      "without-H factorization: io_factorized\n",
      "without-H decay: (1e-06, 1e-06)\n",
      "finite recurrent gradient: True\n",
      "finite feedforward gradient: True\n"
     ]
    }
   ],
   "source": [
    "inputs = brainstate.random.randn(10, 2)\n",
    "targets = jnp.zeros((10, 1))\n",
    "\n",
    "\n",
    "def compile_and_measure(algorithm, recurrent):\n",
    "    model = OSTLModel(recurrent=recurrent)\n",
    "    brainstate.nn.init_all_states(model)\n",
    "    learner = braintrace.compile(model, algorithm, inputs[0])\n",
    "\n",
    "    def step_loss(x, target):\n",
    "        return jnp.mean((learner(x) - target) ** 2)\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=step_loss,\n",
    "        loss_output=\"scalar\",\n",
    "        return_value=True,\n",
    "    )\n",
    "    norm = jnp.sqrt(sum(jnp.sum(g * g) for g in grads.values()))\n",
    "    return learner, loss, norm\n",
    "\n",
    "\n",
    "with_h, recurrent_loss, recurrent_norm = compile_and_measure(\n",
    "    braintrace.OSTLRecurrent, recurrent=True\n",
    ")\n",
    "without_h, feedforward_loss, feedforward_norm = compile_and_measure(\n",
    "    braintrace.OSTLFeedforward, recurrent=False\n",
    ")\n",
    "\n",
    "print(\"with-H recurrence scope:\", with_h.config.recurrence_scope)\n",
    "print(\"without-H factorization:\", without_h.config.trace_factorization)\n",
    "print(\"without-H decay:\", without_h.config.decay)\n",
    "print(\"finite recurrent gradient:\", bool(jnp.isfinite(recurrent_norm)))\n",
    "print(\"finite feedforward gradient:\", bool(jnp.isfinite(feedforward_norm)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a818c552",
   "metadata": {},
   "source": [
    "## 4. Evidence boundary\n",
    "\n",
    "The output verifies that the two public presets compile to different\n",
    "learning-rule coordinates and produce finite gradients on structurally matched\n",
    "models. It does not show that either is exact for a deep network. With-H is\n",
    "gradient-equivalent to BPTT only under the documented block-structured hidden\n",
    "Jacobian conditions; without-H deliberately removes temporal recurrence.\n",
    "\n",
    "## Reference and API\n",
    "\n",
    "- Bohnstingl et al., \"Online Spatio-Temporal Learning in Deep Neural Networks,\"\n",
    "  *IEEE Transactions on Neural Networks and Learning Systems* 34, 8894-8908\n",
    "  (2023), [doi:10.1109/TNNLS.2022.3153985](https://doi.org/10.1109/TNNLS.2022.3153985),\n",
    "  [arXiv:2007.12723](https://arxiv.org/abs/2007.12723).\n",
    "- API: {class}`braintrace.OSTLRecurrent`,\n",
    "  {class}`braintrace.OSTLFeedforward`, {func}`braintrace.compile`.\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
}