{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b4dbbaf7",
   "metadata": {},
   "source": [
    "# SnAp: sparse n-step recurrent influence\n",
    "\n",
    "SnAp-n retains recurrent influence inside an $n$-step dependency neighborhood.\n",
    "A dense recurrent matrix becomes all-to-all almost immediately, so this\n",
    "tutorial uses a sparse ring where increasing $n$ has a visible structural\n",
    "meaning.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3969e3b",
   "metadata": {},
   "source": [
    "## 1. Principle\n",
    "\n",
    "Let $A$ be the directed dependency graph of recurrent hidden positions. SnAp-n\n",
    "keeps trace blocks reachable within the requested order and drops more distant\n",
    "influence. As $n$ grows, the neighborhoods nest; once every position is\n",
    "reachable, the within-group trace saturates. Memory grows with the retained\n",
    "neighborhood width, so $n$ is an accuracy-cost coordinate rather than a generic\n",
    "optimizer setting.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7298a32b",
   "metadata": {},
   "source": [
    "## 2. Applicability and exclusions\n",
    "\n",
    "**Use SnAp when** recurrence is structurally sparse and its position graph can\n",
    "be derived by the compiler.\n",
    "\n",
    "**Do not use a dense recurrence to demonstrate scaling:** its graph diameter is\n",
    "one, so small $n$ already saturates. Unsupported or position-relabeling tails\n",
    "may trigger conservative widening or explicit rejection; inspect diagnostics\n",
    "before interpreting an estimate.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1542fc94",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:15:22.471752Z",
     "iopub.status.busy": "2026-08-07T14:15:22.471752Z",
     "iopub.status.idle": "2026-08-07T14:15:25.852142Z",
     "shell.execute_reply": "2026-08-07T14:15:25.851136Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainevent\n",
    "import brainstate\n",
    "import braintrace\n",
    "import jax.numpy as jnp\n",
    "\n",
    "brainstate.random.seed(41)\n",
    "\n",
    "\n",
    "def ring_structure(n_rec):\n",
    "    row = jnp.arange(n_rec)\n",
    "    mask = jnp.zeros((n_rec, n_rec))\n",
    "    mask = mask.at[row, row].set(1.0)\n",
    "    mask = mask.at[row, (row + 1) % n_rec].set(1.0)\n",
    "    return brainevent.CSR.fromdense(mask), int(mask.sum())\n",
    "\n",
    "\n",
    "class SparseRingRNN(brainstate.nn.Module):\n",
    "    def __init__(self, n_in=2, n_rec=7):\n",
    "        super().__init__()\n",
    "        self.csr, nnz = ring_structure(n_rec)\n",
    "        self.w = brainstate.ParamState(0.25 * brainstate.random.randn(nnz))\n",
    "        self.w_in = brainstate.ParamState(\n",
    "            0.2 * brainstate.random.randn(n_in, n_rec)\n",
    "        )\n",
    "        self.h = brainstate.HiddenState(jnp.zeros(n_rec))\n",
    "\n",
    "    def update(self, x):\n",
    "        recurrent = braintrace.sparse_matmul(\n",
    "            self.h.value, self.w.value, sparse_mat=self.csr\n",
    "        )\n",
    "        self.h.value = jnp.tanh(x @ self.w_in.value + recurrent)\n",
    "        return self.h.value\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd20112f",
   "metadata": {},
   "source": [
    "## 3. Inspect neighborhood growth\n",
    "\n",
    "Each learner receives a fresh model with the same seeded parameter values. The\n",
    "reported `K` is the compiled neighborhood width, not an inferred value from the\n",
    "requested order.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2d538d2e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:15:25.856140Z",
     "iopub.status.busy": "2026-08-07T14:15:25.856140Z",
     "iopub.status.idle": "2026-08-07T14:15:27.948094Z",
     "shell.execute_reply": "2026-08-07T14:15:27.947050Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SnAp-1 (K, trace width, saturated): (1, 1, False)\n",
      "SnAp-2 (K, trace width, saturated): (2, 2, False)\n",
      "SnAp-3 (K, trace width, saturated): (3, 3, False)\n",
      "nested widths: True\n"
     ]
    }
   ],
   "source": [
    "x0 = brainstate.random.randn(2)\n",
    "\n",
    "\n",
    "def compiled_width(order):\n",
    "    with brainstate.random.seed_context(43):\n",
    "        model = SparseRingRNN()\n",
    "    learner = braintrace.compile(model, braintrace.SnAp, x0, n=order)\n",
    "    group = learner.graph.hidden_groups[0]\n",
    "    if group.snap is None:\n",
    "        return 1, group.trace_state_width, False\n",
    "    return group.snap.num_neighbour, group.trace_state_width, group.snap.is_saturated\n",
    "\n",
    "\n",
    "snap1 = compiled_width(1)\n",
    "snap2 = compiled_width(2)\n",
    "snap3 = compiled_width(3)\n",
    "\n",
    "print(\"SnAp-1 (K, trace width, saturated):\", snap1)\n",
    "print(\"SnAp-2 (K, trace width, saturated):\", snap2)\n",
    "print(\"SnAp-3 (K, trace width, saturated):\", snap3)\n",
    "print(\"nested widths:\", snap1[0] <= snap2[0] <= snap3[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4693dd0c",
   "metadata": {},
   "source": [
    "## 4. Execute the public sequence path\n",
    "\n",
    "The next cell checks that a non-saturated order produces a finite online\n",
    "gradient. It does not compare that approximation with BPTT; a finite-window\n",
    "oracle is required for an accuracy claim.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "be3f1bcd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-07T14:15:27.952053Z",
     "iopub.status.busy": "2026-08-07T14:15:27.951055Z",
     "iopub.status.idle": "2026-08-07T14:15:43.454992Z",
     "shell.execute_reply": "2026-08-07T14:15:43.451367Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sequence loss: 0.0816\n",
      "finite nonzero gradient: True\n"
     ]
    }
   ],
   "source": [
    "with brainstate.random.seed_context(43):\n",
    "    model = SparseRingRNN()\n",
    "inputs = brainstate.random.randn(8, 2)\n",
    "targets = jnp.zeros((8, 7))\n",
    "learner = braintrace.compile(model, braintrace.SnAp, inputs[0], n=2)\n",
    "\n",
    "\n",
    "def step_loss(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=step_loss,\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",
    "\n",
    "print(f\"sequence loss: {float(loss):.4f}\")\n",
    "print(\"finite nonzero gradient:\", bool(jnp.isfinite(gradient_norm) & (gradient_norm > 0)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fb7b36f",
   "metadata": {},
   "source": [
    "## 5. Interpretation and limits\n",
    "\n",
    "The ring shows that the compiled neighborhood grows with order instead of\n",
    "saturating immediately. This is structural evidence that the example exercises\n",
    "SnAp's defining mechanism. It is not evidence that order 2 is adequate for a\n",
    "particular task. Accuracy, compiler conservatism, and memory must be evaluated\n",
    "together on the intended sparse architecture.\n",
    "\n",
    "## Reference and API\n",
    "\n",
    "- Menick et al., \"A Practical Sparse Approximation for Real Time Recurrent\n",
    "  Learning,\" ICLR (2021),\n",
    "  [arXiv:2006.07232](https://arxiv.org/abs/2006.07232).\n",
    "- API: {class}`braintrace.SnAp`, {func}`braintrace.sparse_matmul`,\n",
    "  {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
}