{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60001",
   "metadata": {},
   "source": [
    "# Batching Strategies\n",
    "\n",
    "Online learning algorithms need to handle batched data efficiently. In braintrace, there are two main batching strategies:\n",
    "\n",
    "- **Map-based batching** (recommended): Wrap single-sample model logic with `brainstate.nn.Map`, then compile from one complete batched time step.\n",
    "- **Single-sample mode**: Process one sample at a time, without any batching.\n",
    "\n",
    "The choice of strategy affects how model states are initialized and how the online learning algorithm is called.\n",
    "\n",
    "This tutorial walks through each strategy with concrete examples and shows how to build a full training loop using Map-based batching."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60002",
   "metadata": {},
   "source": [
    "## Map-Based Batching (Recommended)\n",
    "\n",
    "Keep the model's update logic single-sample and let `brainstate.nn.Map` manage\n",
    "independent state copies across the batch. Set the mapped learner up explicitly:\n",
    "\n",
    "1. Create exactly one `brainstate.nn.Map(model, init_map_size=B)`.\n",
    "2. Initialize it with `mapped_model.init_all_states()`.\n",
    "3. Construct the online-learning algorithm with the mapped model.\n",
    "4. Compile the ETP graph from one complete batched time step with shape\n",
    "   `(batch_size, n_in)`.\n",
    "\n",
    "Do not pass `mapped_model` to `braintrace.compile(..., vmap=True)`. That setup\n",
    "path owns its batching transformation and would map an already mapped model a\n",
    "second time. When using an explicit `Map`, construct and compile the algorithm\n",
    "directly as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4e5f60003",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import brainstate\n",
    "import braintools\n",
    "import braintrace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4e5f60004",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleGRU(brainstate.nn.Module):\n",
    "    def __init__(self, n_in, n_rec, n_out):\n",
    "        super().__init__()\n",
    "        self.rnn = braintrace.nn.GRUCell(n_in, n_rec)\n",
    "        self.out = braintrace.nn.Linear(n_rec, n_out)\n",
    "\n",
    "    def update(self, x):\n",
    "        return self.out(self.rnn(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4e5f60005",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimpleGRU(10, 64, 5)\n",
    "batch_size = 16\n",
    "example_input = jnp.zeros((batch_size, 10))\n",
    "\n",
    "# Create and initialize exactly one mapped model.\n",
    "mapped_model = brainstate.nn.Map(model, init_map_size=batch_size)\n",
    "mapped_model.init_all_states()\n",
    "\n",
    "# Compile the algorithm directly from the complete batched example input.\n",
    "mapped_algo = braintrace.D_RTRL(mapped_model)\n",
    "mapped_algo.compile_graph(example_input)\n",
    "\n",
    "x_batch = jnp.ones((batch_size, 10))\n",
    "out = mapped_algo(x_batch)\n",
    "print(\"Output shape:\", out.shape)  # (16, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60006",
   "metadata": {},
   "source": [
    "**How it works:**\n",
    "\n",
    "- `brainstate.nn.Map(model, init_map_size=B)` creates the mapped state owner;\n",
    "  `mapped_model.init_all_states()` initializes its independent recurrent states.\n",
    "- The algorithm receives that mapped model and compiles against the complete\n",
    "  batched example input, keeping the batch axis inside the graph.\n",
    "- Each learner call maps the wrapped model over axis 0 while sharing parameter\n",
    "  states and maintaining independent recurrent states.\n",
    "- The Map is created once. It is not passed to another API that would wrap it\n",
    "  again."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60007",
   "metadata": {},
   "source": [
    "## Single-Sample Mode\n",
    "\n",
    "For debugging or situations where batch processing is unnecessary, you can compile and run the algorithm on individual samples directly. No `vmap` or state replication is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4e5f60008",
   "metadata": {},
   "outputs": [],
   "source": [
    "model2 = SimpleGRU(10, 64, 5)\n",
    "\n",
    "# Single-sample mode: omit batch_size so states are created unbatched.\n",
    "algo2 = braintrace.compile(model2, braintrace.D_RTRL, jnp.zeros(10))\n",
    "\n",
    "# Process one sample at a time\n",
    "x_single = jnp.ones(10)\n",
    "out = algo2(x_single)\n",
    "print(\"Single sample output shape:\", out.shape)  # (5,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60009",
   "metadata": {},
   "source": [
    "This mode is straightforward: initialize the model, compile the graph, and call the algorithm. It is useful for step-by-step debugging or when processing a single stream of data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60010",
   "metadata": {},
   "source": [
    "## Multi-Step Data\n",
    "\n",
    "braintrace provides `SingleStepData` and `MultiStepData` wrappers to control how the algorithm processes input along the time dimension.\n",
    "\n",
    "- **`SingleStepData`**: Wraps data for a single time step. The algorithm processes it as one forward pass.\n",
    "- **`MultiStepData`**: Wraps a sequence of time steps. The algorithm internally scans over all steps in the sequence.\n",
    "\n",
    "This is useful when you want to pass an entire sequence to the algorithm and have it handle the temporal loop internally, rather than manually iterating over time steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4e5f60011",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Single-step: process one time step at a time\n",
    "x_single = braintrace.SingleStepData(jnp.ones(10))\n",
    "\n",
    "# Multi-step: process a sequence\n",
    "sequence = jnp.ones((20, 10))  # 20 time steps, 10 features\n",
    "x_multi = braintrace.MultiStepData(sequence)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60012",
   "metadata": {},
   "source": [
    "When a `MultiStepData` object is passed to the algorithm, it will iterate over the first axis (time steps) internally. When a `SingleStepData` object (or a plain array) is passed, the algorithm processes it as a single forward step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60013",
   "metadata": {},
   "source": [
    "## Full Training Loop with Map Batching\n",
    "\n",
    "Below is a complete example that combines Map-based batching with a temporal training loop. The pattern is:\n",
    "\n",
    "1. **Map and initialize** independent model states across the batch.\n",
    "2. **Compile** the algorithm from one batched time step.\n",
    "3. **Drive the sequence** with `etrace_grad`, which walks the time axis and accumulates the per-step online gradients for you — no hand-written scan.\n",
    "4. **Update** parameters with the accumulated gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4e5f60014",
   "metadata": {},
   "outputs": [],
   "source": [
    "@brainstate.transform.jit\n",
    "def train_step(inputs, targets):\n",
    "    \"\"\"inputs: (n_steps, batch_size, n_in), targets: (batch_size,)\"\"\"\n",
    "    def step_loss(inp):\n",
    "        out = mapped_algo(inp)\n",
    "        return jnp.mean((out - targets) ** 2)\n",
    "\n",
    "    # etrace_grad drives the sequence and accumulates per-step online\n",
    "    # gradients. The explicitly mapped model keeps the batch axis inside the\n",
    "    # compiled graph.\n",
    "    return mapped_algo.etrace_grad(\n",
    "        inputs, step_fn=step_loss, reduction='sum'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4e5f60015",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage\n",
    "model = SimpleGRU(10, 64, 5)\n",
    "inputs = jnp.ones((20, 16, 10))  # 20 steps, batch 16, 10 features\n",
    "targets = jnp.zeros((16, 5))\n",
    "\n",
    "mapped_model = brainstate.nn.Map(model, init_map_size=inputs.shape[1])\n",
    "mapped_model.init_all_states()\n",
    "mapped_algo = braintrace.D_RTRL(mapped_model)\n",
    "mapped_algo.compile_graph(inputs[0])\n",
    "\n",
    "grads = train_step(inputs, targets)\n",
    "print(\"Gradient keys:\", list(grads.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60016",
   "metadata": {},
   "source": [
    "**What happens in `train_step`:**\n",
    "\n",
    "1. The setup creates one `brainstate.nn.Map`, initializes it, constructs\n",
    "   `D_RTRL(mapped_model)`, and compiles from the complete batched time step.\n",
    "2. `mapped_algo.etrace_grad` iterates over time, calls `step_fn`, and\n",
    "   accumulates online gradients; `reduction='sum'` accumulates without\n",
    "   dividing.\n",
    "3. The returned gradient keys match `mapped_algo.param_states`. Register those\n",
    "   learner parameter states with the optimizer when applying the gradients."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4e5f60017",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- Create one `brainstate.nn.Map(model, init_map_size=B)` for batched online\n",
    "  learning and call `mapped_model.init_all_states()`.\n",
    "- Construct the algorithm with that mapped model and call\n",
    "  `compile_graph(example_input)` using one complete batched time step.\n",
    "- Never pass an already mapped model to `compile(..., vmap=True)`; doing so\n",
    "  would apply batching twice.\n",
    "- After compilation, drive the time axis with `etrace_grad` (gradients) or\n",
    "  `etrace_evolve` (state/trace only) rather than writing the scan yourself.\n",
    "- For one stream, initialize and compile the original model without `Map`."
   ]
  }
 ],
 "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",
   "nbformat_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}