diff --git a/docs/advanced/batching.ipynb b/docs/advanced/batching.ipynb
index 845b760..195c5f1 100644
--- a/docs/advanced/batching.ipynb
+++ b/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": {
diff --git a/docs/quickstart/concepts.ipynb b/docs/quickstart/concepts.ipynb
index 1bbf058..434ef8c 100644
--- a/docs/quickstart/concepts.ipynb
+++ b/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"
    ]
   },
   {
diff --git a/docs/quickstart/quickstart.ipynb b/docs/quickstart/quickstart.ipynb
index 64d860d..9fe44f2 100644
--- a/docs/quickstart/quickstart.ipynb
+++ b/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)"
diff --git a/docs/tutorials/drtrl.ipynb b/docs/tutorials/drtrl.ipynb
index cf1712b..5239af0 100644
--- a/docs/tutorials/drtrl.ipynb
+++ b/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);"
diff --git a/docs/tutorials/hidden_states.ipynb b/docs/tutorials/hidden_states.ipynb
index c97f311..aefdc61 100644
--- a/docs/tutorials/hidden_states.ipynb
+++ b/docs/tutorials/hidden_states.ipynb
@@ -364,7 +364,7 @@
     "- **Single-sample initialization**: `brainstate.nn.init_all_states(model)` -- state tensors have shape `(M,)`.\n",
     "- **Batched initialization**: `brainstate.nn.init_all_states(model, batch_size=N)` -- state tensors have shape `(N, M)`, where `N` is the batch size. This is used for manual batching.\n",
     "\n",
-    "For automatic batching with `vmap`, you can use `brainstate.transform.vmap_new_states` to initialize per-sample states while keeping the model definition simple."
+    "For automatic batching, wrap the model with `brainstate.nn.Map(model, init_map_size=N)` and call `mapped_model.init_all_states()`. The wrapper keeps the model definition single-sample while managing independent per-sample states."
    ]
   },
   {
@@ -403,7 +403,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "id": "raeddpwnz5o",
    "metadata": {
     "execution": {
@@ -413,29 +413,16 @@
      "shell.execute_reply": "2026-06-26T13:41:44.376744Z"
     }
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "After vmap_new_states, model is ready for automatic batching.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "# --- Automatic batching with vmap_new_states ---\n",
+    "# --- Automatic batching with brainstate.nn.Map ---\n",
     "model = SimpleNeuron(32)\n",
+    "mapped_model = brainstate.nn.Map(model, init_map_size=16)\n",
+    "mapped_model.init_all_states()\n",
     "\n",
-    "@brainstate.transform.vmap_new_states(state_tag='new', axis_size=16)\n",
-    "def init():\n",
-    "    brainstate.nn.init_all_states(model)\n",
-    "\n",
-    "init()\n",
-    "\n",
-    "# After vmap initialization, hidden states are managed per-sample internally.\n",
-    "# The model still \"thinks\" it processes a single sample, but vmap replicates\n",
-    "# the computation across the batch dimension automatically.\n",
-    "print(\"After vmap_new_states, model is ready for automatic batching.\")"
+    "# Map manages independent hidden states for each sample and accepts inputs\n",
+    "# whose leading axis has size 16.\n",
+    "print(\"Mapped model is ready for automatic batching.\")"
    ]
   },
   {
@@ -468,7 +455,20 @@
     }
    },
    "outputs": [],
-   "source": "# Complete example: model -> compile -> inspect graph structure\n\nmodel = SimpleNeuron(8)\n\n# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.\n# Access the compiled graph via learner.graph and the report via learner.report.\nalgo = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(8), batch_size=1)\n\n# Display the discovered graph structure:\n# - Which hidden groups were found\n# - Which weight parameters are associated with each group\nalgo.show_graph()"
+   "source": [
+    "# Complete example: model -> compile -> inspect graph structure\n",
+    "\n",
+    "model = SimpleNeuron(8)\n",
+    "\n",
+    "# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.\n",
+    "# Access the compiled graph via learner.graph and the report via learner.report.\n",
+    "algo = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(8), batch_size=1)\n",
+    "\n",
+    "# Display the discovered graph structure:\n",
+    "# - Which hidden groups were found\n",
+    "# - Which weight parameters are associated with each group\n",
+    "algo.show_graph()"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -496,7 +496,41 @@
     }
    },
    "outputs": [],
-   "source": "# A two-layer recurrent network to demonstrate multi-group discovery\n\nclass TwoLayerRNN(brainstate.nn.Module):\n    \"\"\"Two stacked recurrent layers, each with its own hidden state.\"\"\"\n\n    def __init__(self, in_size, hidden_size, out_size):\n        super().__init__()\n        # Layer 1\n        self.w1_in = brainstate.ParamState(brainstate.random.randn(in_size, hidden_size) * 0.01)\n        self.w1_rec = brainstate.ParamState(brainstate.random.randn(hidden_size, hidden_size) * 0.01)\n        self.h1 = brainstate.HiddenState(jnp.zeros(hidden_size))\n\n        # Layer 2\n        self.w2_in = brainstate.ParamState(brainstate.random.randn(hidden_size, out_size) * 0.01)\n        self.w2_rec = brainstate.ParamState(brainstate.random.randn(out_size, out_size) * 0.01)\n        self.h2 = brainstate.HiddenState(jnp.zeros(out_size))\n\n    def update(self, x):\n        # Layer 1: x feeds in, h1 recurs\n        self.h1.value = jax.nn.tanh(\n            x @ self.w1_in.value + braintrace.matmul(self.h1.value, self.w1_rec.value)\n        )\n        # Layer 2: h1 feeds in, h2 recurs\n        self.h2.value = jax.nn.tanh(\n            self.h1.value @ self.w2_in.value + braintrace.matmul(self.h2.value, self.w2_rec.value)\n        )\n        return self.h2.value\n\n\nmodel_2layer = TwoLayerRNN(in_size=10, hidden_size=16, out_size=8)\n\nalgo_2layer = braintrace.compile(model_2layer, braintrace.D_RTRL, jnp.zeros(10), batch_size=1)\nalgo_2layer.show_graph()"
+   "source": [
+    "# A two-layer recurrent network to demonstrate multi-group discovery\n",
+    "\n",
+    "class TwoLayerRNN(brainstate.nn.Module):\n",
+    "    \"\"\"Two stacked recurrent layers, each with its own hidden state.\"\"\"\n",
+    "\n",
+    "    def __init__(self, in_size, hidden_size, out_size):\n",
+    "        super().__init__()\n",
+    "        # Layer 1\n",
+    "        self.w1_in = brainstate.ParamState(brainstate.random.randn(in_size, hidden_size) * 0.01)\n",
+    "        self.w1_rec = brainstate.ParamState(brainstate.random.randn(hidden_size, hidden_size) * 0.01)\n",
+    "        self.h1 = brainstate.HiddenState(jnp.zeros(hidden_size))\n",
+    "\n",
+    "        # Layer 2\n",
+    "        self.w2_in = brainstate.ParamState(brainstate.random.randn(hidden_size, out_size) * 0.01)\n",
+    "        self.w2_rec = brainstate.ParamState(brainstate.random.randn(out_size, out_size) * 0.01)\n",
+    "        self.h2 = brainstate.HiddenState(jnp.zeros(out_size))\n",
+    "\n",
+    "    def update(self, x):\n",
+    "        # Layer 1: x feeds in, h1 recurs\n",
+    "        self.h1.value = jax.nn.tanh(\n",
+    "            x @ self.w1_in.value + braintrace.matmul(self.h1.value, self.w1_rec.value)\n",
+    "        )\n",
+    "        # Layer 2: h1 feeds in, h2 recurs\n",
+    "        self.h2.value = jax.nn.tanh(\n",
+    "            self.h1.value @ self.w2_in.value + braintrace.matmul(self.h2.value, self.w2_rec.value)\n",
+    "        )\n",
+    "        return self.h2.value\n",
+    "\n",
+    "\n",
+    "model_2layer = TwoLayerRNN(in_size=10, hidden_size=16, out_size=8)\n",
+    "\n",
+    "algo_2layer = braintrace.compile(model_2layer, braintrace.D_RTRL, jnp.zeros(10), batch_size=1)\n",
+    "algo_2layer.show_graph()"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -526,7 +560,7 @@
     "1. **Automatic discovery**: The compiler traces the model's Jaxpr and automatically identifies which states are recurrent. No manual annotation of hidden states is needed -- just use `brainstate`'s state classes.\n",
     "2. **Grouping**: Related hidden states are grouped together for efficient Jacobian computation. `HiddenGroupState` and `HiddenTreeState` explicitly declare a group; separate `HiddenState` variables are grouped by data flow analysis.\n",
     "3. **Operation-based selection**: Whether a weight participates in online learning depends on the operation used (`braintrace.matmul` vs. regular `@`), not on the parameter class.\n",
-    "4. **Flexible initialization**: Use `init_all_states` for single-sample or manual batching, and `vmap_new_states` for automatic batching."
+    "4. **Flexible initialization**: Use `init_all_states` for single-sample or manual batching, and `brainstate.nn.Map` plus `mapped_model.init_all_states()` for automatic batching."
    ]
   }
  ],
@@ -551,4 +585,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
\ No newline at end of file
+}
diff --git a/docs/tutorials/neural_network_layers.ipynb b/docs/tutorials/neural_network_layers.ipynb
index cbfbb00..ebab467 100644
--- a/docs/tutorials/neural_network_layers.ipynb
+++ b/docs/tutorials/neural_network_layers.ipynb
@@ -33,8 +33,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import warnings\n",
-    "\n",
     "import brainstate\n",
     "import jax.numpy as jnp\n",
     "\n",
@@ -93,16 +91,11 @@
     "\n",
     "model = TinySequenceModel()\n",
     "sample = jnp.ones(1)\n",
-    "with warnings.catch_warnings():\n",
-    "    warnings.filterwarnings(\n",
-    "        \"ignore\",\n",
-    "        message=r\"ETP primitive etp_mv.*has no connected hidden states.*\",\n",
-    "    )\n",
-    "    learner = braintrace.compile(\n",
-    "        model,\n",
-    "        braintrace.D_RTRL,\n",
-    "        sample,\n",
-    "    )\n"
+    "learner = braintrace.compile(\n",
+    "    model,\n",
+    "    braintrace.D_RTRL,\n",
+    "    sample,\n",
+    ")\n"
    ]
   },
   {
diff --git a/docs/tutorials/pp_prop.ipynb b/docs/tutorials/pp_prop.ipynb
index 35b7f07..54738c1 100644
--- a/docs/tutorials/pp_prop.ipynb
+++ b/docs/tutorials/pp_prop.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,19 +101,15 @@
     "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.pp_prop,\n",
-    "        inputs[0],\n",
-    "        batch_size=1,\n",
-    "        decay_or_rank=0.9,\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.pp_prop,\n",
+    "    inputs[0],\n",
+    "    batch_size=1,\n",
+    "    decay_or_rank=0.9,\n",
+    ")\n",
     "weights = model.states(brainstate.ParamState)\n",
     "optimizer = braintools.optim.SGD(lr=0.08)\n",
     "optimizer.register_trainable_weights(weights);"
diff --git a/docs/tutorials/rnn_online_learning.ipynb b/docs/tutorials/rnn_online_learning.ipynb
index f90bbc5..a3192a9 100644
--- a/docs/tutorials/rnn_online_learning.ipynb
+++ b/docs/tutorials/rnn_online_learning.ipynb
@@ -55,7 +55,6 @@
     "import braintools\n",
     "import braintrace\n",
     "import matplotlib.pyplot as plt\n",
-    "import warnings\n",
     "\n",
     "brainstate.random.seed(17)"
    ]
@@ -170,16 +169,27 @@
    "source": [
     "## 4. Online Training with D-RTRL\n",
     "\n",
-    "D-RTRL (Diagonal Real-Time Recurrent Learning) is an online learning algorithm provided by `braintrace`. Unlike BPTT, which requires storing the entire computation graph across all time steps, D-RTRL computes approximate gradients incrementally at each time step using **eligibility traces**. It is not generally gradient-equivalent to BPTT outside the assumptions of its diagonal Jacobian approximation.\n",
+    "D-RTRL (Diagonal Real-Time Recurrent Learning) is an online learning algorithm\n",
+    "provided by `braintrace`. Unlike BPTT, which requires storing the entire\n",
+    "computation graph across all time steps, D-RTRL computes approximate gradients\n",
+    "incrementally using **eligibility traces**. It is not generally\n",
+    "gradient-equivalent to BPTT outside the assumptions of its diagonal Jacobian\n",
+    "approximation.\n",
     "\n",
     "The key steps in the online training loop are:\n",
     "\n",
-    "1. **Compile the model**: `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initialises hidden states, compiles the eligibility-trace graph, and returns a vmapped learner — all in one call.\n",
-    "2. **Warm-up phase**: Use `learner.etrace_evolve(...)` to advance hidden states and eligibility traces without computing a loss gradient.\n",
-    "3. **Learning phase**: Use `learner.etrace_grad(..., step_fn=step_loss)` to drive the remaining sequence and accumulate online gradients.\n",
-    "4. **Parameter update**: After processing the full sequence, apply the accumulated gradients to update the parameters.\n",
-    "\n",
-    "The `D_RTRL` class wraps the model and handles the eligibility trace bookkeeping automatically."
+    "1. **Map once**: Create `brainstate.nn.Map(model, init_map_size=B)` and call\n",
+    "   `mapped_model.init_all_states()`.\n",
+    "2. **Compile directly**: Construct `braintrace.D_RTRL(mapped_model)` and compile\n",
+    "   from one complete batched time step.\n",
+    "3. **Warm up**: Use `learner.etrace_evolve(...)` to advance hidden states and\n",
+    "   eligibility traces without computing a loss gradient.\n",
+    "4. **Learn**: Use `learner.etrace_grad(..., step_fn=step_loss)` to accumulate\n",
+    "   online gradients, then update the parameters.\n",
+    "\n",
+    "An already mapped model must not be passed to\n",
+    "`braintrace.compile(..., vmap=True)`, because that would apply a second mapping\n",
+    "layer."
    ]
   },
   {
@@ -200,22 +210,20 @@
     "    \"\"\"Train one GRU with D-RTRL over precomputed copying batches.\"\"\"\n",
     "    brainstate.random.seed(21)\n",
     "    model = GRUNet(10, 64, 10)\n",
-    "    opt = braintools.optim.Adam(lr)\n",
-    "    weights = model.states().subset(brainstate.ParamState)\n",
-    "    opt.register_trainable_weights(weights)\n",
-    "\n",
     "    batch_size = input_batches.shape[2]\n",
-    "    with warnings.catch_warnings():\n",
-    "        warnings.filterwarnings(\"ignore\", message=r\"ETP primitive .*\")\n",
-    "        learner = braintrace.compile(\n",
-    "            model, braintrace.D_RTRL, input_batches[0, 0],\n",
-    "            batch_size=batch_size, vmap=True,\n",
-    "        )\n",
+    "\n",
+    "    mapped_model = brainstate.nn.Map(model, init_map_size=batch_size)\n",
+    "    mapped_model.init_all_states()\n",
+    "    learner = braintrace.D_RTRL(mapped_model)\n",
+    "    learner.compile_graph(input_batches[0])\n",
+    "\n",
+    "    opt = braintools.optim.Adam(lr)\n",
+    "    opt.register_trainable_weights(learner.param_states)\n",
     "\n",
     "    @brainstate.transform.jit\n",
     "    def train_step(inputs, targets):\n",
-    "        brainstate.nn.reset_all_states(model, batch_size=batch_size)\n",
-    "        learner.reset_state(batch_size=batch_size)\n",
+    "        brainstate.nn.reset_all_states(mapped_model)\n",
+    "        learner.reset_state()\n",
     "\n",
     "        # The loss for ONE step. `etrace_grad` owns the loop; this owns the\n",
     "        # model call, so multi-head models and regularizers need no special\n",
@@ -279,23 +287,19 @@
     "\n",
     "    @brainstate.transform.jit\n",
     "    def train_step(inputs, targets):\n",
-    "        @brainstate.transform.vmap_new_states(\n",
-    "            state_tag=\"new\", axis_size=inputs.shape[1]\n",
+    "        mapped_model = brainstate.nn.Map(\n",
+    "            model, init_map_size=inputs.shape[1]\n",
     "        )\n",
-    "        def init():\n",
-    "            brainstate.nn.init_all_states(model)\n",
-    "\n",
-    "        init()\n",
-    "        vmapped_model = brainstate.nn.Vmap(model, vmap_states=\"new\")\n",
+    "        mapped_model.init_all_states()\n",
     "\n",
     "        def run_step(inp, tar):\n",
-    "            out = vmapped_model(inp)\n",
+    "            out = mapped_model(inp)\n",
     "            loss = braintools.metric.softmax_cross_entropy_with_integer_labels(out, tar).mean()\n",
     "            return out, loss\n",
     "\n",
     "        def bptt_forward():\n",
     "            n_sim = time_lag + 10\n",
-    "            brainstate.transform.for_loop(vmapped_model, inputs[:n_sim])\n",
+    "            brainstate.transform.for_loop(mapped_model, inputs[:n_sim])\n",
     "            outs, losses = brainstate.transform.for_loop(run_step, inputs[n_sim:], targets)\n",
     "            return losses.mean(), outs\n",
     "\n",
@@ -442,18 +446,23 @@
    "source": [
     "## 8. Summary\n",
     "\n",
-    "In this tutorial, we demonstrated how to use `braintrace` for online learning of a GRU network on the copying task.\n",
+    "In this tutorial, we demonstrated online learning of a GRU on the copying task.\n",
     "\n",
     "**Key takeaways:**\n",
     "\n",
-    "- **D-RTRL** provides approximate online gradients with `O(B * theta)` complexity, where `B` is the batch size and `theta` is the number of parameters. Unlike BPTT, it does not need to store the full unrolled computation graph.\n",
-    "- The online training loop uses `braintrace.compile` to set up the algorithm. A single call to `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initialises hidden states, compiles the eligibility-trace graph, and returns a vmapped learner ready for batched training.\n",
-    "- Online learning uses `learner.etrace_evolve` for gradient-free sequence prefixes and `learner.etrace_grad` for sequence objectives; both compose with `brainstate.transform.jit`.\n",
-    "- `braintrace` is particularly effective for RNN models with gating mechanisms (GRU, LSTM), where the internal dynamics naturally support eligibility trace propagation.\n",
+    "- **D-RTRL** provides approximate online gradients with `O(B * theta)`\n",
+    "  complexity, where `B` is the batch size and `theta` is the number of\n",
+    "  parameters. Unlike BPTT, it does not store the full unrolled graph.\n",
+    "- Batched online learning creates one `brainstate.nn.Map`, initializes it, and\n",
+    "  passes it directly to `braintrace.D_RTRL` before `compile_graph` is called on\n",
+    "  a complete batched time step.\n",
+    "- Do not pass an already mapped model to `compile(..., vmap=True)`.\n",
+    "- Use `learner.etrace_evolve` for gradient-free prefixes and\n",
+    "  `learner.etrace_grad` for sequence objectives.\n",
     "\n",
     "For more details, see:\n",
     "- [Key Concepts](../quickstart/concepts.ipynb) for the theoretical background.\n",
-    "- [SNN Online Learning](./snn_online_learning.ipynb) for applying the same approach to spiking neural networks."
+    "- [SNN Online Learning](./snn_online_learning.ipynb) for spiking networks."
    ]
   }
  ],
diff --git a/docs/tutorials/snn_online_learning.ipynb b/docs/tutorials/snn_online_learning.ipynb
index 21586a9..fdc5e36 100644
--- a/docs/tutorials/snn_online_learning.ipynb
+++ b/docs/tutorials/snn_online_learning.ipynb
@@ -62,7 +62,6 @@
     "import brainunit as u\n",
     "import brainpy.state\n",
     "import matplotlib.pyplot as plt\n",
-    "import warnings\n",
     "\n",
     "brainstate.random.seed(31)"
    ]
@@ -198,16 +197,24 @@
    "source": [
     "## 3. Training with ES-D-RTRL\n",
     "\n",
-    "We set up online learning using `braintrace.compile` with `braintrace.pp_prop` (also exposed as `braintrace.ES_D_RTRL` and `braintrace.IODimVjpAlgorithm`). The `decay_or_rank` parameter controls the trace approximation:\n",
+    "We use `braintrace.pp_prop` (also exposed as `braintrace.ES_D_RTRL` and\n",
+    "`braintrace.IODimVjpAlgorithm`). The `decay_or_rank` parameter controls the\n",
+    "trace approximation:\n",
     "\n",
-    "* **`decay_or_rank=float` in (0, 1)** -- exponentially-smoothed trace. Larger values retain a longer history but may be less stable; this short teaching task uses `0.5`. Memory cost: `O(B * (I + O))` per layer.\n",
-    "* **`decay_or_rank=int >= 1`** -- an alternative parameterization of the same decay, using `decay = (rank - 1) / (rank + 1)`. It does not allocate multiple rank factors, so memory remains `O(B * (I + O))`.\n",
+    "* **`decay_or_rank=float` in (0, 1)** -- exponentially smoothed trace. Larger\n",
+    "  values retain a longer history; this short teaching task uses `0.5`.\n",
+    "* **`decay_or_rank=int >= 1`** -- an alternative parameterization using\n",
+    "  `decay = (rank - 1) / (rank + 1)`.\n",
     "\n",
-    "Use the float form to set the decay directly, or the integer form to select the corresponding decay through the documented conversion. Neither form creates an independent rank-versus-memory trade-off.\n",
+    "For batched learning, create one `brainstate.nn.Map`, initialize it, construct\n",
+    "`braintrace.pp_prop(mapped_model, decay_or_rank=0.5)`, and compile from one\n",
+    "complete batched time step. Do not pass the mapped model to\n",
+    "`braintrace.compile(..., vmap=True)`, because that would map it a second time.\n",
     "\n",
-    "`braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` initialises per-sample states, builds the ETP graph, and returns a vmapped learner — no separate `init_all_states`, `compile_graph`, or `Vmap` calls are needed.\n",
-    "\n",
-    "`braintrace.D_RTRL` is the alternative parameter-dimensional estimator. It stores parameter-shaped eligibility traces and can use substantially more memory; neither estimator is generally gradient-equivalent to BPTT outside its documented assumptions."
+    "`braintrace.D_RTRL` is the alternative parameter-dimensional estimator. It\n",
+    "stores parameter-shaped eligibility traces and can use substantially more\n",
+    "memory; neither estimator is generally gradient-equivalent to BPTT outside its\n",
+    "documented assumptions."
    ]
   },
   {
@@ -249,25 +256,22 @@
     "    with brainstate.environ.context(dt=1. * u.ms):\n",
     "        brainstate.random.seed(37)\n",
     "        model = LIF_SNN(n_in, n_rec, n_out)\n",
+    "\n",
+    "        mapped_model = brainstate.nn.Map(model, init_map_size=batch_size)\n",
+    "        mapped_model.init_all_states()\n",
+    "        learner = braintrace.pp_prop(mapped_model, decay_or_rank=0.5)\n",
+    "        learner.compile_graph(input_batches[0])\n",
+    "\n",
     "        opt = braintools.optim.Adam(lr)\n",
-    "        weights = model.states(brainstate.ParamState)\n",
-    "        opt.register_trainable_weights(weights)\n",
-    "\n",
-    "        with warnings.catch_warnings():\n",
-    "            warnings.filterwarnings(\"ignore\", message=r\"ETP primitive .*\")\n",
-    "            learner = braintrace.compile(\n",
-    "                model, braintrace.pp_prop, input_batches[0, 0],\n",
-    "                batch_size=batch_size, vmap=True, decay_or_rank=0.5,\n",
-    "            )\n",
+    "        opt.register_trainable_weights(learner.param_states)\n",
     "\n",
     "        @brainstate.transform.jit\n",
     "        def train_step(inputs, targets):\n",
-    "            brainstate.nn.reset_all_states(model, batch_size=batch_size)\n",
-    "            learner.reset_state(batch_size=batch_size)\n",
+    "            brainstate.nn.reset_all_states(mapped_model)\n",
+    "            learner.reset_state()\n",
     "\n",
     "            # One step's loss. `etrace_grad` owns the loop; `step_fn` owns the\n",
-    "            # model call and hands the step's logits back as aux, so they come\n",
-    "            # out of the driver stacked over time for the accuracy readout.\n",
+    "            # model call and returns logits as auxiliary output.\n",
     "            def step_loss(inp):\n",
     "                output = learner(inp)\n",
     "                loss = braintools.metric.softmax_cross_entropy_with_integer_labels(\n",
@@ -275,10 +279,6 @@
     "                ).mean()\n",
     "                return loss, output\n",
     "\n",
-    "            # One call slices the sequence, differentiates each step's loss\n",
-    "            # online, and accumulates the per-step gradients. `reduction='sum'`\n",
-    "            # keeps the accumulated scale this example's learning rate was\n",
-    "            # tuned at.\n",
     "            grads, step_losses, outputs = learner.etrace_grad(\n",
     "                inputs, step_fn=step_loss, has_aux=True,\n",
     "                reduction='sum', return_value=True,\n",
@@ -424,17 +424,22 @@
    "source": [
     "## 5. Summary\n",
     "\n",
-    "In this tutorial, we demonstrated how to train a spiking neural network with online learning using BrainTrace. Here are the key takeaways:\n",
-    "\n",
-    "1. **Model Construction**: Use `braintrace.nn.Linear` and `braintrace.nn.LeakyRateReadout` for layers that should participate in online learning (ETP-aware). Combine them with spiking neuron models from `brainpy.state` (e.g., `LIF`).\n",
-    "\n",
-    "2. **Online Learning Setup**: Use `learner = braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` to initialise states, compile the ETP graph, and return a vmapped learner in one call. Then use `learner.etrace_grad(...)` to drive the sequence and accumulate online gradients.\n",
+    "This tutorial demonstrated online learning for a spiking neural network.\n",
     "\n",
-    "3. **Scalability**: ES-D-RTRL achieves O(B(I+O)) memory complexity, making it practical for large spiking networks. The `decay_or_rank` parameter controls the trace approximation quality.\n",
+    "1. **Model construction**: Use ETP-aware BrainTrace layers for parameters that\n",
+    "   should participate in online learning.\n",
+    "2. **Mapped setup**: Create and initialize one `brainstate.nn.Map`, pass it\n",
+    "   directly to `braintrace.pp_prop`, and compile from a complete batched time\n",
+    "   step.\n",
+    "3. **No duplicate mapping**: Never pass an already mapped model to\n",
+    "   `compile(..., vmap=True)`.\n",
+    "4. **Sequence training**: Use `learner.etrace_grad(...)` to drive the sequence\n",
+    "   and accumulate online gradients.\n",
     "\n",
-    "4. **Batching**: `braintrace.compile(..., batch_size=B, vmap=True)` handles per-sample state initialisation and vmapped execution automatically.\n",
+    "The separate `braintools`/`saiunit` physical-unit compatibility issue is not\n",
+    "handled by this documentation change.\n",
     "\n",
-    "For more advanced topics, including training on real neuromorphic datasets (N-MNIST) and comparing online learning with BPTT, see:\n",
+    "For more advanced topics, see:\n",
     "- [pp_prop algorithm tutorial](pp_prop.ipynb)\n",
     "- [Key Concepts](../quickstart/concepts.ipynb)"
    ]