Fix warning visibility and Map-based state initialization

docs/advanced/batching.ipynb

@@ -9,19 +9,35 @@
     "\n",
     "Online learning algorithms need to handle batched data efficiently. In braintrace, there are two main batching strategies:\n",
     "\n",
-    "- **Vmap-based batching** (recommended): Compile the computation graph for a single sample, then use `vmap` to automatically vectorize across the batch dimension.\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 vmap-based batching."
+    "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": "## Vmap-Based Batching (Recommended)\n\nThe recommended approach is to compile the online learning graph for a **single sample**, then leverage JAX's `vmap` to parallelize across the batch. `braintrace.compile(..., batch_size=B, vmap=True)` does this in one call:\n\n1. It initialises per-sample states for `batch_size` samples.\n2. It compiles the ETP graph from a single **batched** time step (shape `(batch_size, n_in)`); the batch axis (axis 0) is stripped internally to recover the per-sample example.\n3. It wraps the algorithm with `brainstate.nn.Vmap` for parallel execution.\n4. It returns the vmapped learner, ready to call on batched inputs.\n\nThis pattern keeps the model definition simple (single-sample logic) while gaining efficient batch parallelism automatically."
+   "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",
@@ -60,13 +76,40 @@
    "id": "a1b2c3d4e5f60005",
    "metadata": {},
    "outputs": [],
-   "source": "model = SimpleGRU(10, 64, 5)\nbatch_size = 16\n\n# braintrace.compile with vmap=True:\n#   - initialises per-sample hidden states (batch_size independent copies)\n#   - compiles the ETP graph from a single batched time step: axis 0 is the batch\n#     axis, which compile strips internally to recover the per-sample example\n#   - wraps the result in brainstate.nn.Vmap for parallel execution\nalgo_vmapped = braintrace.compile(\n    model, braintrace.D_RTRL, jnp.zeros((batch_size, 10)),\n    batch_size=batch_size, vmap=True,\n)\n\n# Run on batched input — the returned learner handles the batch axis transparently\nx_batch = jnp.ones((batch_size, 10))\nout = algo_vmapped(x_batch)\nprint(\"Output shape:\", out.shape)  # (16, 5)"
+   "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- `braintrace.compile(..., batch_size=B, vmap=True)` internally runs `brainstate.transform.vmap_new_states` to create `B` independent copies of all model states (tagged `'new'`), then compiles the ETP graph on a single-sample input, and finally wraps the learner in `brainstate.nn.Vmap(algo, vmap_states='new')` for parallel execution.\n- Each call to `algo_vmapped(x_batch)` automatically splits the batch input across the per-sample states, runs the forward pass independently for each sample, and stacks the outputs.\n- The model itself only ever sees single-sample inputs — all batch handling is transparent."
+   "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",
@@ -84,7 +127,17 @@
    "id": "a1b2c3d4e5f60008",
    "metadata": {},
    "outputs": [],
-   "source": "model2 = SimpleGRU(10, 64, 5)\n\n# Single-sample mode: omit batch_size so states are created unbatched.\nalgo2 = braintrace.compile(model2, braintrace.D_RTRL, jnp.zeros(10))\n\n# Process one sample at a time\nx_single = jnp.ones(10)\nout = algo2(x_single)\nprint(\"Single sample output shape:\", out.shape)  # (5,)"
+   "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",
@@ -137,12 +190,12 @@
    "id": "a1b2c3d4e5f60013",
    "metadata": {},
    "source": [
-    "## Full Training Loop with Vmap Batching\n",
+    "## Full Training Loop with Map Batching\n",
     "\n",
-    "Below is a complete example that combines vmap-based batching with a temporal training loop. The pattern is:\n",
+    "Below is a complete example that combines Map-based batching with a temporal training loop. The pattern is:\n",
     "\n",
-    "1. **Initialize** model states and compile the graph for a single sample.\n",
-    "2. **Vmap** the algorithm across the batch dimension.\n",
+    "1. **Map and initialize** independent model states across the batch.\n",
+    "2. **Compile** the algorithm from one batched time step.\n",
     "3. **Scan** over time steps, accumulating gradients at each step.\n",
     "4. **Update** parameters with the accumulated gradients."
    ]
@@ -157,22 +210,16 @@
     "@brainstate.transform.jit\n",
     "def train_step(inputs, targets):\n",
     "    \"\"\"inputs: (n_steps, batch_size, n_in), targets: (batch_size,)\"\"\"\n",
-    "    # braintrace.compile with vmap=True replaces the manual init + compile_graph + Vmap pattern.\n",
-    "    # Pass the batched single time step inputs[0] (shape (batch_size, n_in)); compile strips\n",
-    "    # axis 0 internally to recover the per-sample example — do NOT pass inputs[0, 0].\n",
-    "    vmapped_algo = braintrace.compile(\n",
-    "        model, braintrace.D_RTRL, inputs[0],\n",
-    "        batch_size=inputs.shape[1], vmap=True,\n",
-    "    )\n",
-    "\n",
     "    def step_loss(inp):\n",
-    "        out = vmapped_algo(inp)\n",
+    "        out = mapped_algo(inp)\n",
     "        return jnp.mean((out - targets) ** 2)\n",
     "\n",
-    "    # etrace_grad drives the whole sequence and accumulates the per-step online\n",
-    "    # gradients. The vmapped learner carries the same driver methods as an\n",
-    "    # unbatched one, so nothing here changes when you switch batching strategy.\n",
-    "    return vmapped_algo.etrace_grad(inputs, step_fn=step_loss, reduction='sum')"
+    "    # 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",
+    "    )"
    ]
   },
   {
@@ -186,6 +233,12 @@
     "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()))"
    ]
@@ -197,16 +250,31 @@
    "source": [
     "**What happens in `train_step`:**\n",
     "\n",
-    "1. `braintrace.compile(model, braintrace.D_RTRL, inputs[0], batch_size=B, vmap=True)` initialises per-sample hidden states, compiles the computation graph using a single time step's batched input (`inputs[0]`, shape `(batch_size, n_in)`), and wraps the result in a `braintrace.ETraceVmap` (a `brainstate.nn.Vmap` that also carries the sequence drivers) for batch-parallel execution — all in one call.\n",
-    "2. `vmapped_algo.etrace_grad` iterates over the time dimension (`inputs` has shape `(n_steps, batch_size, n_in)`). At each step it calls `step_fn` for the loss and takes its online gradient, then accumulates; `reduction='sum'` accumulates without dividing. The vmapped learner carries the same drivers as an unbatched one, so this line is identical in both batching strategies.\n",
-    "3. The returned `grads` dictionary can be passed to an optimizer (e.g., `braintools.optim.Adam`) for a parameter update."
+    "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- **`braintrace.compile(..., batch_size=B, vmap=True)` is the recommended way** to set up batched online learning. It replaces the three-step manual pattern (`init_all_states` + `compile_graph` + `Vmap`) with a single call.\n- The workflow is: **`braintrace.compile` (with `vmap=True`) → scan over time steps → accumulate gradients → update parameters**.\n- `SingleStepData` and `MultiStepData` control whether the algorithm processes one time step or scans over an entire sequence internally.\n- For non-batched (single-sample) use, call `braintrace.compile` without `vmap=True`; states are initialised for a single sample."
+   "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",
+    "- Scan over time with `etrace_grad` or `etrace_evolve` after compilation.\n",
+    "- For one stream, initialize and compile the original model without `Map`."
+   ]
   }
  ],
  "metadata": {

docs/quickstart/concepts.ipynb

@@ -115,10 +115,7 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import brainstate\n",
-    "import braintrace\n",
-    "import warnings\n",
-    "\n",
-    "warnings.filterwarnings(\"ignore\", message=r\"ETP primitive .*\")"
+    "import braintrace"
    ]
   },
   {

docs/quickstart/quickstart.ipynb

@@ -80,8 +80,6 @@
     }
    ],
    "source": [
-    "import warnings\n",
-    "\n",
     "import brainstate\n",
     "import braintools\n",
     "import braintrace\n",
@@ -107,18 +105,14 @@
     "targets = 0.7 * inputs + 0.2\n",
     "\n",
     "# Compile once. inputs[0] is one batched time step with shape (1, 1).\n",
-    "with warnings.catch_warnings():\n",
-    "    # The readout does not feed recurrent state, so it is non-temporal.\n",
-    "    warnings.filterwarnings(\n",
-    "        \"ignore\",\n",
-    "        message=r\"ETP primitive etp_mm.*has no connected hidden states.*\",\n",
-    "    )\n",
-    "    learner = braintrace.compile(\n",
-    "        model,\n",
-    "        braintrace.D_RTRL,\n",
-    "        inputs[0],\n",
-    "        batch_size=1,\n",
-    "    )\n",
+    "# The compiler reports that the readout is non-temporal because it does not\n",
+    "# feed a recurrent state.\n",
+    "learner = braintrace.compile(\n",
+    "    model,\n",
+    "    braintrace.D_RTRL,\n",
+    "    inputs[0],\n",
+    "    batch_size=1,\n",
+    ")\n",
     "weights = model.states(brainstate.ParamState)\n",
     "optimizer = braintools.optim.SGD(lr=0.08)\n",
     "optimizer.register_trainable_weights(weights)"

docs/tutorials/drtrl.ipynb

@@ -59,7 +59,6 @@
     "import braintrace\n",
     "import jax.numpy as jnp\n",
     "import matplotlib.pyplot as plt\n",
-    "import warnings\n",
     "\n",
     "brainstate.random.seed(7)"
    ]
@@ -102,18 +101,14 @@
     "inputs = jnp.linspace(-1.0, 1.0, 12).reshape(12, 1, 1)\n",
     "targets = 0.7 * inputs + 0.2\n",
     "\n",
-    "with warnings.catch_warnings():\n",
-    "    # The Linear readout is intentionally non-temporal; Section 5 explains why.\n",
-    "    warnings.filterwarnings(\n",
-    "        \"ignore\",\n",
-    "        message=r\"ETP primitive etp_mm.*has no connected hidden states.*\",\n",
-    "    )\n",
-    "    learner = braintrace.compile(\n",
-    "        model,\n",
-    "        braintrace.D_RTRL,\n",
-    "        inputs[0],\n",
-    "        batch_size=1,\n",
-    "    )\n",
+    "# The Linear readout is intentionally non-temporal; Section 5 explains\n",
+    "# the compiler diagnostic emitted for it.\n",
+    "learner = braintrace.compile(\n",
+    "    model,\n",
+    "    braintrace.D_RTRL,\n",
+    "    inputs[0],\n",
+    "    batch_size=1,\n",
+    ")\n",
     "weights = model.states(brainstate.ParamState)\n",
     "optimizer = braintools.optim.SGD(lr=0.08)\n",
     "optimizer.register_trainable_weights(weights);"