diff --git a/Cargo.lock b/Cargo.lock
index 864aa8e5..040fcbdc 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3728,6 +3728,7 @@ dependencies = [
  "prost",
  "pyo3",
  "pyo3-build-config",
+ "sqlparser",
  "tokio",
  "tonic",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 482e10ee..300b5c9e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,6 +32,7 @@ prost = { version = "0.14", optional = true }
 # enables `pyo3/extension-module` through pyproject.toml for wheel builds; it
 # must stay disabled for ordinary Cargo test binaries so they link libpython.
 pyo3 = { version = "0.29.0", features = ["abi3-py310"] }
+sqlparser = { version = "0.62", features = ["visitor"] }
 tokio = { version = "1.46.1", features = ["macros", "rt", "sync", "time"] }
 tonic = { version = "0.14", features = ["transport"], optional = true }
 
diff --git a/accelerators/einfold_mlx/build.py b/accelerators/einfold_mlx/build.py
new file mode 100644
index 00000000..7991034f
--- /dev/null
+++ b/accelerators/einfold_mlx/build.py
@@ -0,0 +1,97 @@
+"""Build the pinned EinFold dense-executor hook into this Python environment.
+
+EinFold (https://github.com/xqlsystems/einfold) is a DataFusion tensor
+optimizer. Its released Python API has no external kernel interface. This
+script builds and installs a wheel from the pinned feature commit in an
+isolated checkout. Requires git, a Rust toolchain, maturin,
+and uv. Nothing is installed into the system Python.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+_REVISION = "0b2f6db892c3cbf9df2c3d775fbcf4c5654df093"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--source", type=Path, help="reuse a clean pinned source checkout"
+    )
+    args = parser.parse_args()
+    if sys.prefix == sys.base_prefix:
+        parser.error(
+            "run this builder with the checkout's virtual-environment Python"
+        )
+    for command in ("git", "cargo", "uv"):
+        if shutil.which(command) is None:
+            parser.error(f"{command} must be on PATH")
+    with tempfile.TemporaryDirectory(prefix="xql-einfold-") as temporary:
+        work = Path(temporary)
+        source = args.source or work / "source"
+        if args.source is None:
+            subprocess.run(
+                [
+                    "git",
+                    "clone",
+                    "https://github.com/cromwellian/einfold.git",
+                    str(source),
+                ],
+                check=True,
+            )
+            subprocess.run(
+                ["git", "checkout", _REVISION], cwd=source, check=True
+            )
+        revision = subprocess.check_output(
+            ["git", "rev-parse", "HEAD"], cwd=source, text=True
+        ).strip()
+        if revision != _REVISION:
+            parser.error(
+                f"expected EinFold revision {_REVISION}, got {revision}"
+            )
+        if subprocess.check_output(
+            ["git", "status", "--porcelain"], cwd=source
+        ):
+            parser.error(
+                "the source checkout must be clean"
+            )
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "maturin",
+                "build",
+                "--release",
+                "--locked",
+                "--out",
+                str(work / "wheels"),
+            ],
+            cwd=source / "python",
+            env={**os.environ, "PYO3_PYTHON": sys.executable},
+            check=True,
+        )
+        (wheel,) = (work / "wheels").glob("*.whl")
+        subprocess.run(
+            [
+                "uv",
+                "pip",
+                "install",
+                "--python",
+                sys.executable,
+                "--force-reinstall",
+                "--no-deps",
+                str(wheel),
+            ],
+            check=True,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py
index 44f5c3cc..ce4edbb1 100644
--- a/benchmarks/ml/01_attention.py
+++ b/benchmarks/ml/01_attention.py
@@ -231,7 +231,15 @@ def main() -> None:
                 rows = report.time(
                     f"forward ({spelling})", host, n, host.query, FORWARD
                 )
-                report.check(assert_close("forward", matrix(rows), want))
+                report.check(
+                    assert_close(
+                        "forward",
+                        matrix(rows),
+                        want,
+                        rtol=host.rtol,
+                        atol=host.atol,
+                    )
+                )
                 report.note("operators", host.operators(FORWARD))
 
                 step = f"gradient ({spelling})"
@@ -242,7 +250,13 @@ def main() -> None:
                 if rows is None:
                     continue
                 errors = [
-                    assert_close(f"d/d{name}", matrix(rows, f"d{name}"), ref)
+                    assert_close(
+                        f"d/d{name}",
+                        matrix(rows, f"d{name}"),
+                        ref,
+                        rtol=host.rtol,
+                        atol=host.atol,
+                    )
                     for name, ref in zip("qkv", want_grads)
                 ]
                 report.check(max(errors))
@@ -303,7 +317,11 @@ def check_nadaraya_watson(host_names: Sequence[str], report: Report) -> None:
             NADARAYA_WATSON,
         )
         got = xql.to_dataset(rows, dims=["p"])["v"].values
-        report.check(assert_close("nadaraya-watson", got, want))
+        report.check(
+            assert_close(
+                "nadaraya-watson", got, want, rtol=host.rtol, atol=host.atol
+            )
+        )
 
 
 def check_einfold_fuses(host_names: Sequence[str]) -> None:
diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py
index 3dc477c4..d8e4ad1b 100644
--- a/benchmarks/ml/02_nanogpt.py
+++ b/benchmarks/ml/02_nanogpt.py
@@ -85,6 +85,7 @@
 from __future__ import annotations
 
 import pathlib
+import time
 import urllib.request
 
 import jax
@@ -167,11 +168,33 @@ def layer_norm(x: str, out: str, gain: str, layer: int | None = None) -> str:
 GELU_SCALE = np.sqrt(2 / np.pi)
 
 
-def attention(layer: int, x: str, out: str) -> str:
+def attention(
+    layer: int, x: str, out: str, *, dropout=False, grouped_softmax=False
+) -> str:
     """nanoGPT's ``CausalSelfAttention``: each position of ``x`` attends to
     itself and the positions before it, in each head, and ``out`` is the
     heads' outputs, projected back to the residual stream."""
     l = f"l{layer}"
+    mask = (
+        f"JOIN {l}_adrop d ON a.seq=d.seq AND a.h=d.h AND a.i=d.i AND a.j=d.j"
+        if dropout
+        else ""
+    )
+    scale = " * d.scale" if dropout else ""
+    projected = f"{out}_raw" if dropout else out
+    weights = f"""
+{l}_weights AS (
+  SELECT seq, h, i, j,
+         exp(s - MAX(s) OVER (PARTITION BY seq, h, i)) AS p
+  FROM {l}_scores),"""
+    if grouped_softmax:
+        weights = f"""
+{l}_maximum AS (
+  SELECT seq, h, i, ddx_stop_gradient(MAX(s)) AS m FROM {l}_scores GROUP BY seq, h, i),
+{l}_weights AS (
+  SELECT a.seq, a.h, a.i, a.j, exp(a.s - m.m) AS p
+  FROM {l}_scores a JOIN {l}_maximum m
+  ON a.seq=m.seq AND a.h=m.h AND a.i=m.i),"""
     projections = "".join(
         f"""
 {l}_{name} AS (
@@ -188,30 +211,37 @@ def attention(layer: int, x: str, out: str) -> str:
   FROM {l}_q q JOIN {l}_k k ON q.seq = k.seq AND q.h = k.h AND q.c = k.c
   WHERE q.pos >= k.pos
   GROUP BY q.seq, q.h, q.pos, k.pos),
-{l}_weights AS (
-  -- The max only keeps exp() from overflowing: it cancels out of p / z.
-  SELECT seq, h, i, j,
-         exp(s - MAX(s) OVER (PARTITION BY seq, h, i)) AS p
-  FROM {l}_scores),
+{weights}
 {l}_norm AS (
   SELECT seq, h, i, SUM(p) AS z FROM {l}_weights GROUP BY seq, h, i),
 {l}_heads AS (
-  SELECT a.seq, a.i AS pos, a.h, v.c, SUM(a.p / n.z * v.val) AS val
+  SELECT a.seq, a.i AS pos, a.h, v.c, SUM(a.p / n.z * v.val{scale}) AS val
   FROM {l}_weights a
   JOIN {l}_v v ON a.seq = v.seq AND a.h = v.h AND a.j = v.pos
   JOIN {l}_norm n ON a.seq = n.seq AND a.h = n.h AND a.i = n.i
+  {mask}
   GROUP BY a.seq, a.i, a.h, v.c),
-{out} AS (
+{projected} AS (
   SELECT a.seq, a.pos, p.e, SUM(a.val * p.w) AS val
   FROM {l}_heads a JOIN wo p ON a.h = p.h AND a.c = p.c
   WHERE p.layer = {layer}
-  GROUP BY a.seq, a.pos, p.e),"""
+  GROUP BY a.seq, a.pos, p.e),""" + (
+        _dropout(projected, out, f"{l}_rdrop") if dropout else ""
+    )
 
 
-def mlp(layer: int, x: str, out: str) -> str:
+def mlp(
+    layer: int, x: str, out: str, *, exact_gelu=False, dropout=False
+) -> str:
     """nanoGPT's ``MLP``: each position of ``x`` through a hidden layer
     four times as wide, and GELU, back to the residual stream as ``out``."""
     l = f"l{layer}"
+    activation = (
+        "0.5 * val * (1 + erf(val / sqrt(2.0)))"
+        if exact_gelu
+        else f"0.5 * val * (1 + tanh({GELU_SCALE} * (val + 0.044715 * val * val * val)))"
+    )
+    projected = f"{out}_raw" if dropout else out
     return f"""
 {l}_fc AS (
   SELECT x.seq, x.pos, p.f, SUM(x.val * p.w) AS val
@@ -220,14 +250,22 @@ def mlp(layer: int, x: str, out: str) -> str:
   GROUP BY x.seq, x.pos, p.f),
 {l}_hidden AS (
   SELECT seq, pos, f,
-         0.5 * val * (1 + tanh({GELU_SCALE} * (val + 0.044715 * val * val * val)))
-           AS val
+         {activation} AS val
   FROM {l}_fc),
-{out} AS (
+{projected} AS (
   SELECT x.seq, x.pos, p.e, SUM(x.val * p.w) AS val
   FROM {l}_hidden x JOIN w_proj p ON x.f = p.f
   WHERE p.layer = {layer}
-  GROUP BY x.seq, x.pos, p.e),"""
+  GROUP BY x.seq, x.pos, p.e),""" + (
+        _dropout(projected, out, f"{l}_mdrop") if dropout else ""
+    )
+
+
+def _dropout(x, out, mask):
+    return f"""
+{out} AS (
+  SELECT x.seq, x.pos, x.e, x.val * d.scale AS val
+  FROM {x} x JOIN {mask} d ON x.seq=d.seq AND x.pos=d.pos AND x.e=d.e),"""
 
 
 def residual(x: str, y: str, out: str) -> str:
@@ -238,38 +276,71 @@ def residual(x: str, y: str, out: str) -> str:
   FROM {x} x JOIN {y} y ON x.seq = y.seq AND x.pos = y.pos AND x.e = y.e),"""
 
 
-def block(layer: int, x: str, out: str) -> str:
+def block(
+    layer: int,
+    x: str,
+    out: str,
+    *,
+    exact_gelu=False,
+    dropout=False,
+    grouped_softmax=False,
+) -> str:
     """nanoGPT's ``Block``: ``mid = x + attention(ln_1(x))``, then
     ``out = mid + mlp(ln_2(mid))``."""
     l = f"l{layer}"
     return "".join(
         [
             layer_norm(x, f"{l}_ln1", "ln_1", layer),
-            attention(layer, f"{l}_ln1", f"{l}_attn"),
+            attention(
+                layer,
+                f"{l}_ln1",
+                f"{l}_attn",
+                dropout=dropout,
+                grouped_softmax=grouped_softmax,
+            ),
             residual(x, f"{l}_attn", f"{l}_mid"),
             layer_norm(f"{l}_mid", f"{l}_ln2", "ln_2", layer),
-            mlp(layer, f"{l}_ln2", f"{l}_mlp"),
+            mlp(
+                layer,
+                f"{l}_ln2",
+                f"{l}_mlp",
+                exact_gelu=exact_gelu,
+                dropout=dropout,
+            ),
             residual(f"{l}_mid", f"{l}_mlp", out),
         ]
     )
 
 
-def model() -> str:
+def model(
+    *, exact_gelu=False, dropout=False, view="logits", grouped_softmax=False
+) -> str:
     """The model, as the view ``logits``: for each position of each
     sequence in ``batch``, a score for every character that could come
     next. A DataFusion view reads the tables that existed when it was
     created (DuckDB's looks them up by name), so it is recreated whenever
     the batch or the parameters are."""
     blocks = "".join(
-        block(layer, f"x{layer}", f"x{layer + 1}") for layer in range(LAYERS)
+        block(
+            layer,
+            f"x{layer}",
+            f"x{layer + 1}",
+            exact_gelu=exact_gelu,
+            dropout=dropout,
+            grouped_softmax=grouped_softmax,
+        )
+        for layer in range(LAYERS)
     )
+    embedded = "embed_raw" if dropout else "x0"
+    embed_dropout = _dropout(embedded, "x0", "embed_drop") if dropout else ""
     return f"""
-CREATE OR REPLACE VIEW logits AS
-WITH x0 AS (
+CREATE OR REPLACE VIEW {view} AS
+WITH {embedded} AS (
   SELECT b.seq, b.pos, te.e, te.w + pe.w AS val
   FROM batch b
   JOIN wte te ON b.tok = te.tok
   JOIN wpe pe ON b.pos = pe.pos AND te.e = pe.e),
+{embed_dropout}
 {blocks}
 {layer_norm(f"x{LAYERS}", "final", "ln_f").rstrip(",")}
 SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val
@@ -460,15 +531,7 @@ def init_params(vocab: int, rng) -> dict[str, np.ndarray]:
     """nanoGPT's initialization, from GPT-2: weights normal with standard
     deviation 0.02, or 0.02/√(2 × layers) for the two projections that add
     to the residual stream; layer norm gains 1."""
-    sizes = {
-        "tok": vocab,
-        "pos": CONTEXT,
-        "layer": LAYERS,
-        "e": EMBED,
-        "h": HEADS,
-        "c": HEAD,
-        "f": HIDDEN,
-    }
+    sizes = _parameter_sizes(vocab)
     params = {}
     for name, dims in PARAMS.items():
         shape = [sizes[d] for d in dims]
@@ -481,6 +544,18 @@ def init_params(vocab: int, rng) -> dict[str, np.ndarray]:
     return params
 
 
+def _parameter_sizes(vocab):
+    return {
+        "tok": vocab,
+        "pos": CONTEXT,
+        "layer": LAYERS,
+        "e": EMBED,
+        "h": HEADS,
+        "c": HEAD,
+        "f": HIDDEN,
+    }
+
+
 def as_table(name: str, w: np.ndarray, m=None, v=None) -> xr.Dataset:
     """A parameter as a table: its weights ``w``, and Adam's moments ``m``
     and ``v`` when they are given."""
@@ -527,19 +602,62 @@ def main():
     global LAYERS, EMBED, HEADS, HEAD, HIDDEN, CONTEXT
     parser = cli(__doc__.splitlines()[0], sizes=[4], size_name="batch")
     parser.add_argument("--steps", type=int, default=5, help="training steps")
+    parser.add_argument(
+        "--learning-rate",
+        type=float,
+        default=LEARNING_RATE,
+        help="SGD learning rate (AdamW uses its schedule)",
+    )
+    parser.add_argument(
+        "--preset",
+        choices=["tiny", "demo", "shakespeare", "gpt2"],
+        default="tiny",
+        help="model dimensions; explicit size flags override them",
+    )
+    parser.add_argument(
+        "--vocab-size",
+        type=int,
+        help="pad vocabulary (GPT-2 preset defaults to 50257; character inputs)",
+    )
+    parser.add_argument(
+        "--forward-only",
+        action="store_true",
+        help="benchmark loss without parameter updates",
+    )
     parser.add_argument(
         "--optimizer",
         choices=["sgd", "adamw"],
         default="sgd",
         help="SGD at a fixed learning rate (default), or nanoGPT's AdamW",
     )
-    parser.add_argument("--layers", type=int, default=LAYERS, help="n_layer")
-    parser.add_argument("--embed", type=int, default=EMBED, help="n_embd")
-    parser.add_argument("--heads", type=int, default=HEADS, help="n_head")
-    parser.add_argument(
-        "--context", type=int, default=CONTEXT, help="block_size"
-    )
+    parser.add_argument("--layers", type=int, help="n_layer")
+    parser.add_argument("--embed", type=int, help="n_embd")
+    parser.add_argument("--heads", type=int, help="n_head")
+    parser.add_argument("--context", type=int, help="block_size")
     args = parser.parse_args()
+    if not np.isfinite(args.learning_rate) or args.learning_rate <= 0:
+        parser.error("--learning-rate must be finite and positive")
+    preset = {
+        "tiny": (1, 16, 2, 16),
+        "demo": (1, 128, 4, 64),
+        "shakespeare": (6, 384, 6, 256),
+        "gpt2": (12, 768, 12, 1024),
+    }[args.preset]
+    for name, value in zip(("layers", "embed", "heads", "context"), preset):
+        if getattr(args, name) is None:
+            setattr(args, name, value)
+    if (
+        min(
+            args.steps,
+            args.layers,
+            args.embed,
+            args.heads,
+            args.context,
+            *args.sizes,
+        )
+        <= 0
+    ):
+        parser.error("steps, model dimensions and batch sizes must be positive")
     if args.embed % args.heads:
         parser.error("--embed must be a multiple of --heads")
     LAYERS, EMBED, HEADS, CONTEXT = (
@@ -554,14 +672,43 @@ def main():
     report = Report("nanogpt", "batch", args.reps)
     text = shakespeare()
     vocab = len(set(text))
+    if args.vocab_size is None and args.preset == "gpt2":
+        args.vocab_size = 50257
+    if args.vocab_size is not None:
+        if args.vocab_size < vocab:
+            parser.error(
+                "--vocab-size must cover the Shakespeare character vocabulary"
+            )
+        vocab = args.vocab_size
     lines = len(model().splitlines()) + len(loss_statement().splitlines())
+    sizes = _parameter_sizes(vocab)
     count = sum(
-        v.size for v in init_params(vocab, np.random.default_rng(0)).values()
+        int(np.prod([sizes[d] for d in dims])) for dims in PARAMS.values()
     )
+    metadata = {
+        "layers": LAYERS,
+        "embed": EMBED,
+        "heads": HEADS,
+        "context": CONTEXT,
+        "vocab": vocab,
+        "parameters": count,
+        "optimizer": args.optimizer,
+        "sgd_learning_rate": args.learning_rate if not adamw else None,
+        "forward_only": args.forward_only,
+    }
     print(
         f"  {LAYERS}-layer GPT, vocab {vocab}, {count:,} parameters: "
         f"the loss is {lines} lines of SQL"
     )
+    print(
+        f"  steps={args.steps}, optimizer={args.optimizer}, "
+        f"context={CONTEXT}, heads={HEADS}, embed={EMBED}, "
+        f"forward_only={args.forward_only}"
+    )
+    if vocab != len(set(text)):
+        print(
+            "  Padded vocabulary: inputs are character IDs, without GPT-2 BPE tokenization"
+        )
 
     for size in args.sizes:
         init = init_params(vocab, np.random.default_rng(0))
@@ -572,18 +719,34 @@ def main():
         )
 
         for host in hosts(args.hosts):
+            load_seconds = 0.0
+            model_seconds = 0.0
+
+            def load_table(name, values, host=host):
+                nonlocal load_seconds
+                start = time.perf_counter()
+                host.load(name, values)
+                load_seconds += time.perf_counter() - start
+
             reference = dict(init)
-            moments = {
-                n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items()
-            }
+            moments = (
+                {
+                    n: (np.zeros_like(w), np.zeros_like(w))
+                    for n, w in init.items()
+                }
+                if adamw
+                else {}
+            )
             for name, values in init.items():
                 state = moments[name] if adamw else ()
-                host.load(name, as_table(name, values, *state))
-            host.load("schedule", lr)
+                load_table(name, as_table(name, values, *state))
+            load_table("schedule", lr)
 
             for step, batch in data:
-                host.load("batch", batch)
+                load_table("batch", batch)
+                start = time.perf_counter()
                 host.execute(model())
+                model_seconds += time.perf_counter() - start
                 tok, nxt = batch["tok"].values, batch["next"].values
 
                 rows = report.time(
@@ -595,10 +758,28 @@ def main():
                 )
                 want = jax_loss(reference, tok, nxt)
                 report.check(
-                    assert_close("loss", rows["l"][0].as_py(), want, rtol=1e-9)
+                    assert_close(
+                        "loss",
+                        rows["l"][0].as_py(),
+                        want,
+                        rtol=host.rtol,
+                        atol=host.atol,
+                    )
                 )
                 report.note("loss", round(float(want), 6))
                 report.note("operators", host.operators(loss_statement()))
+                report.note("rtol", f"{host.rtol:.0e}")
+                report.note("atol", f"{host.atol:.0e}")
+                for key, value in metadata.items():
+                    report.note(key, value)
+                for key, value in host.measurements().items():
+                    report.note(key, value)
+                if args.forward_only:
+                    report.note("load_seconds", load_seconds)
+                    report.note("model_seconds", model_seconds)
+                    report._write_csv(args.csv)
+                    continue
+                report._write_csv(args.csv)
 
                 train = f"train ({args.optimizer}), step {step}"
                 if not host.differentiates:
@@ -614,7 +795,7 @@ def main():
                         host,
                         size,
                         host.query_all,
-                        step_statements(LEARNING_RATE),
+                        step_statements(args.learning_rate),
                     )
                 if updated is None:
                     # Every later step starts from this one's weights.
@@ -626,7 +807,7 @@ def main():
                     )
                 else:
                     reference = {
-                        name: w - LEARNING_RATE * grads[name]
+                        name: w - args.learning_rate * grads[name]
                         for name, w in reference.items()
                     }
 
@@ -641,12 +822,21 @@ def main():
                                 f"{name}.{column}, step {step}",
                                 as_array(name, rows, column),
                                 values,
+                                rtol=host.rtol,
+                                atol=host.atol,
                             )
                         )
-                    host.load(
+                    load_table(
                         name, xql.to_dataset(rows, dims=list(PARAMS[name]))
                     )
                 report.check(max(errors))
+                report.note("load_seconds", load_seconds)
+                report.note("model_seconds", model_seconds)
+                for key, value in metadata.items():
+                    report.note(key, value)
+                for key, value in host.measurements().items():
+                    report.note(key, value)
+                report._write_csv(args.csv)
 
     report.finish(args.csv)
 
diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py
index dcb044e8..3888c8fb 100644
--- a/benchmarks/ml/_bench.py
+++ b/benchmarks/ml/_bench.py
@@ -24,7 +24,7 @@
 from dataclasses import dataclass, field
 
 import numpy as np
-from _hosts import HOSTS, Host
+from _hosts import DEFAULT_HOSTS, HOSTS, Host
 
 
 def cli(description: str, sizes: Sequence[int], size_name: str = "n"):
@@ -33,9 +33,9 @@ def cli(description: str, sizes: Sequence[int], size_name: str = "n"):
     parser.add_argument(
         "--hosts",
         nargs="+",
-        default=list(HOSTS),
+        default=list(DEFAULT_HOSTS),
         choices=list(HOSTS),
-        help="engines to run on (default: all)",
+        help="engines to run on (default: datafusion einfold duckdb)",
     )
     parser.add_argument(
         "--sizes",
@@ -156,6 +156,15 @@ def note(self, key: str, value) -> None:
         """Attach a measurement besides time to the last row."""
         self.rows[-1][key] = value
 
+    def _write_csv(self, path):
+        """Checkpoint completed checks during long model runs."""
+        if path:
+            columns = list(dict.fromkeys(k for row in self.rows for k in row))
+            with open(path, "w", newline="") as f:
+                writer = csv.DictWriter(f, fieldnames=columns)
+                writer.writeheader()
+                writer.writerows(self.rows)
+
     def finish(self, csv_path: str | None = None) -> None:
         """Print the table, and write it as CSV if asked. Exits with an
         error if a host refused a step it wasn't expected to."""
@@ -166,10 +175,7 @@ def finish(self, csv_path: str | None = None) -> None:
             cells = [_cell(row.get(c)) for c in columns]
             print(f"| {' | '.join(cells)} |")
         if csv_path:
-            with open(csv_path, "w", newline="") as f:
-                writer = csv.DictWriter(f, fieldnames=columns)
-                writer.writeheader()
-                writer.writerows(self.rows)
+            self._write_csv(csv_path)
             print(f"\nwrote {csv_path}", file=sys.stderr)
         if self.unexpected:
             sys.exit(
diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py
index b002b79b..1d6ea982 100644
--- a/benchmarks/ml/_hosts.py
+++ b/benchmarks/ml/_hosts.py
@@ -27,13 +27,21 @@
 A gradient that ddx (or the engine under it) refuses to plan raises one
 of the host's ``refusals``, and the demos report that step with the reason.
 
-Adding a host (DuckDB with a GPU extension, say) is one class here and
-one entry in ``HOSTS``.
+Optional Apple silicon hosts use MLX (https://github.com/ml-explore/mlx),
+Apple's array framework, through xarray-sql's native dense-kernel bridge.
+They pair prepared queries with reusable parameter buffers. Forced variants
+lower kernel thresholds to measure small operations on each backend.
+
+Adding a host is one class here and one entry in ``HOSTS``. ``DEFAULT_HOSTS``
+keeps optional device dependencies out of ordinary runs.
 """
 
 from __future__ import annotations
 
 import json
+import os
+import re
+from collections import Counter
 from collections.abc import Sequence
 
 import duckdb
@@ -45,6 +53,21 @@
 
 import xarray_sql as xql
 
+try:
+    from xarray_sql.mlx import MLX
+except ImportError:
+    MLX = None  # Optional Apple silicon dependency and Metal device.
+
+try:
+    from xarray_sql.shared import SharedExecutor
+except ImportError:
+    SharedExecutor = None
+
+try:
+    from xarray_sql.dense import DenseExecutor
+except ImportError:
+    DenseExecutor = None
+
 
 class Host:
     """An engine that runs the demos' SQL. Subclasses fill in the rest."""
@@ -52,6 +75,11 @@ class Host:
     name = ""
     differentiates = False
     refusals: tuple[type[Exception], ...] = ()
+    rtol = 1e-9
+    atol = 1e-9
+
+    def measurements(self) -> dict:
+        return {}
 
     def load(self, name: str, data: xr.Dataset, **options) -> None:
         raise NotImplementedError
@@ -114,7 +142,262 @@ class EinFold(DataFusion):
     def __init__(self):
         super().__init__()
         self.einfold = EinFoldExtension()
-        self.ctx.with_extensions(self.einfold)
+        self.ctx = self.ctx.with_extensions(self.einfold)
+
+    def operators(self, sql):
+        return None  # Native physical stages cross the FFI as opaque nodes.
+
+
+class DataFusionPrepared(DataFusion):
+    """Cached SQL plans and ddx programs over reusable parameter buffers.
+
+    Data is copied on first registration so in-place training updates never
+    mutate the caller's initial weights or the batches used by other hosts.
+    A schema/coordinate change refreshes the tables, views and query caches.
+    """
+
+    name = "datafusion_prepared"
+
+    def __init__(self):
+        super().__init__()
+        for kind in ("", "_join", "_topk", "_aggregate"):
+            self.ctx.sql(
+                f"SET datafusion.optimizer.enable{kind}_dynamic_filter_pushdown = false"
+            ).collect()
+        self._tables: set[str] = set()
+        self._queries: dict = {}
+        self._programs: dict = {}
+        self._last_view: str | None = None
+        self._counts: Counter[str] = Counter()
+        self.accelerator = None
+
+    def load(self, name, data, **options):
+        if name in self._tables:
+            try:
+                self.ctx._write({name: data})
+            except ValueError:
+                pass
+            else:
+                self._counts["inplace_updates"] += 1
+                return
+        self._queries.clear()
+        self._programs.clear()
+        self._last_view = None
+        super().load(name, data.copy(deep=True), **options)
+        self._tables.add(name)
+        self._counts["table_registers"] += 1
+
+    def execute(self, sql):
+        # The GPT recreates this same view after loading each batch. When
+        # tables retain their buffers, its existing plan reads current values.
+        view = sql.lstrip().upper().startswith("CREATE OR REPLACE VIEW ")
+        if view and sql == self._last_view:
+            return
+        super().execute(sql)
+        self._queries.clear()
+        self._programs.clear()
+        self._last_view = sql if view else None
+
+    def query(self, sql):
+        if xql.ad.calls_ad(sql):
+            return self.query_all([sql])[0]
+        if sql not in self._queries:
+            self._queries[sql] = self.ctx.prepare(sql)
+            self._counts["query_preparations"] += 1
+        return self._queries[sql]()
+
+    def query_all(self, sqls):
+        key = tuple(sqls)
+        if key not in self._programs:
+            self._programs[key] = xql.ad.prepare(self.ctx, key)
+            self._counts["gradient_preparations"] += 1
+        return self._programs[key]()
+
+    def measurements(self):
+        values = dict(self._counts)
+        if self.accelerator is not None:
+            counters = self.accelerator.statistics()
+            values.update(
+                {
+                    key: counters.get(key, 0)
+                    for key in (
+                        "cpu_matmul",
+                        "gpu_matmul",
+                        "gpu_attention",
+                        "declined",
+                    )
+                }
+            )
+        return values
+
+    def operators(self, sql):
+        return None if self.accelerator is not None else super().operators(sql)
+
+
+class CPUOptimized(DataFusionPrepared):
+    """Prepared execution plus the float64 CPU dense-kernel bridge."""
+
+    name = "cpu_optimized"
+    precision = "float64"
+    force_kernels = False
+
+    def __init__(self):
+        super().__init__()
+        if MLX is None:
+            raise ImportError(
+                "this host needs the Apple silicon MLX extra and native hook"
+            )
+        options = {"min_ops": 0, "min_gpu_ops": 0} if self.force_kernels else {}
+        self.accelerator = MLX(precision=self.precision, **options)
+        self.ctx = self.ctx.with_extensions(self.accelerator)
+
+
+class CPUForced(CPUOptimized):
+    """Force eligible small products through the float64 CPU backend."""
+
+    name = "cpu_forced"
+    force_kernels = True
+
+
+class GPUOptimized(CPUOptimized):
+    """Prepared execution with the measured automatic GPU thresholds."""
+
+    name = "gpu_optimized"
+    precision = "float32"
+    rtol = 3e-5
+    atol = 3e-5
+
+
+class GPUForced(GPUOptimized):
+    """Force eligible operations onto MLX to expose small-kernel overhead."""
+
+    name = "gpu_forced"
+    force_kernels = True
+
+
+class SharedEinFold(DataFusionPrepared):
+    """Shared forward CTEs and reusable ddx physical plans through EinFold."""
+
+    name = "shared_einfold"
+    precision = None
+
+    def __init__(self):
+        super().__init__()
+        if SharedExecutor is None:
+            raise ImportError("rebuild xarray-sql for shared SQL execution")
+        extension = EinFoldExtension()
+        if self.precision is not None:
+            if MLX is None:
+                raise ImportError(
+                    "shared CPU/GPU kernels require the MLX extra"
+                )
+            extension = self.accelerator = MLX(
+                precision=self.precision,
+                min_ops=0,
+                min_gpu_ops=0,
+                max_working_bytes=4 << 30,
+            )
+        self.ctx = self.ctx.with_extensions(extension)
+        self.executor = SharedExecutor(self.ctx, max_key_cache_bytes=4 << 30)
+
+    def load(self, name, data, **options):
+        # Large stacked weights remain lazy Arrow streams. One layer per
+        # partition lets a block's constant layer filter prune the others.
+        if "layer" in data.dims and "chunks" not in options:
+            options["chunks"] = {"layer": 1}
+        super().load(name, data, **options)
+
+    def execute(self, sql):
+        view = re.fullmatch(
+            r"\s*CREATE OR REPLACE VIEW (\w+) AS\s*(.*)",
+            sql,
+            re.IGNORECASE | re.DOTALL,
+        )
+        if view:
+            name, query = view.groups()
+            if sql != self._last_view:
+                # Preserve the score/softmax chain as one EinFold region.
+                self.executor.inline_ctes = frozenset(
+                    re.findall(
+                        r"\b(l\d+_(?:scores|weights|norm))\s+AS",
+                        query,
+                        re.IGNORECASE,
+                    )
+                )
+                self.executor.register_view(name, query)
+                self._queries.clear()
+                self._programs.clear()
+                self._last_view = sql
+            return
+        super().execute(sql)
+
+    def query(self, sql):
+        if xql.ad.calls_ad(sql):
+            return self.query_all([sql])[0]
+        if sql not in self._queries:
+            self._queries[sql] = self.executor.prepare(sql)
+            self._counts["query_preparations"] += 1
+        return self._queries[sql]()
+
+    def query_all(self, sqls):
+        key = tuple(sqls)
+        if key not in self._programs:
+            self._programs[key] = self.executor.prepare_all(sqls)
+            self._counts["gradient_preparations"] += 1
+        return self._programs[key]()
+
+    def operators(self, sql):
+        # The FFI planner exposes each stage as an opaque node. Its node
+        # count cannot be compared with ordinary DataFusion's operators.
+        return None
+
+    def measurements(self):
+        return super().measurements() | self.executor.statistics()
+
+
+class SharedCPU(SharedEinFold):
+    """Shared execution with double-precision dense CPU kernels."""
+
+    name = "shared_cpu"
+    precision = "float64"
+
+
+class SharedGPU(SharedEinFold):
+    """Shared execution with explicitly enabled MLX float32 GPU kernels."""
+
+    name = "shared_gpu"
+    precision = "float32"
+    rtol = 3e-5
+    atol = 3e-5
+
+
+class DenseCPU(SharedEinFold):
+    """Compact double-precision arrays across SQL and ddx stages."""
+
+    name = "dense_cpu"
+    dense_precision = "float64"
+
+    def __init__(self):
+        super().__init__()
+        if DenseExecutor is None:
+            raise ImportError("install EinFold's optional dense program module")
+        self.executor = DenseExecutor(
+            self.ctx,
+            precision=self.dense_precision,
+            fallback=False,
+            max_cached_bytes=int(
+                os.environ.get("XQL_DENSE_CACHE_BYTES", str(8 << 30))
+            ),
+        )
+
+
+class DenseGPU(DenseCPU):
+    """Compact MLX arrays across SQL and ddx stages, with float32 arithmetic."""
+
+    name = "dense_gpu"
+    dense_precision = "float32"
+    rtol = 3e-5
+    atol = 3e-5
 
 
 class DuckDB(Host):
@@ -159,4 +442,22 @@ def _count(node: dict) -> int:
     return 1 + sum(_count(child) for child in node.get("children", []))
 
 
-HOSTS = {host.name: host for host in (DataFusion, EinFold, DuckDB)}
+DEFAULT_HOSTS = (DataFusion.name, EinFold.name, DuckDB.name)
+HOSTS = {
+    host.name: host
+    for host in (
+        DataFusion,
+        EinFold,
+        DuckDB,
+        DataFusionPrepared,
+        CPUOptimized,
+        CPUForced,
+        GPUOptimized,
+        GPUForced,
+        SharedEinFold,
+        SharedCPU,
+        SharedGPU,
+        DenseCPU,
+        DenseGPU,
+    )
+}
diff --git a/pyproject.toml b/pyproject.toml
index d7b32771..49242e07 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,6 +38,12 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+# Apple silicon GPU kernels need the native hook built by
+# accelerators/einfold_mlx/build.py in addition to these dependencies.
+mlx = [
+  "einfold[datafusion]>=0.2,<0.3; sys_platform == 'darwin' and platform_machine == 'arm64'",
+  "mlx>=0.32.4,<0.33; sys_platform == 'darwin' and platform_machine == 'arm64'",
+]
 # Lazy, chunked results from `to_dataset(chunks=...)` need a chunk manager.
 dask = [
   "dask>=2024.8.0",
@@ -143,7 +149,7 @@ dev = [
     "ddxdb[duckdb]>=0.4.6,<0.5; sys_platform != 'emscripten'",
     # TODO(xqlsystems/einfold#118): install einfold on Windows too, once
     # it publishes a Windows wheel; until then it would build from source.
-    "einfold[datafusion]>=0.1.0rc2,<0.2; sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "einfold[datafusion]>=0.1.0rc2,<0.3; sys_platform != 'emscripten' and sys_platform != 'win32'",
     # jaxlib publishes no wheels for Intel macOS, so the demos skip there.
     "jax; sys_platform != 'darwin' or platform_machine != 'x86_64'",
 ]
diff --git a/src/lib.rs b/src/lib.rs
index 003d8707..2fd796b3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -43,6 +43,7 @@
 
 #[cfg(feature = "flight")]
 mod flight;
+mod shared;
 
 use std::collections::{HashMap, HashSet};
 use std::ffi::CString;
@@ -1656,6 +1657,8 @@ impl LazyArrowStreamTable {
 #[pymodule]
 fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<LazyArrowStreamTable>()?;
+    m.add_function(wrap_pyfunction!(shared::shared_ctes, m)?)?;
+    m.add_class::<shared::SortedArrowTable>()?;
     #[cfg(feature = "flight")]
     m.add_class::<flight::FlightSqlServer>()?;
     Ok(())
diff --git a/src/shared.rs b/src/shared.rs
new file mode 100644
index 00000000..160a81f9
--- /dev/null
+++ b/src/shared.rs
@@ -0,0 +1,137 @@
+//! Parse an explicitly staged SELECT with the same SQL parser as DataFusion.
+
+use std::ops::ControlFlow;
+use std::sync::Arc;
+
+use arrow::array::RecordBatch;
+use arrow::pyarrow::FromPyArrow;
+use datafusion::catalog::TableProvider;
+use datafusion::common::Column;
+use datafusion::datasource::MemTable;
+use datafusion::logical_expr::Expr as DFExpr;
+use datafusion_ffi::table_provider::FFI_TableProvider;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+use sqlparser::ast::{Expr, ObjectName, Query, Statement, Visit, Visitor};
+use sqlparser::dialect::GenericDialect;
+use sqlparser::parser::Parser;
+
+#[derive(Default)]
+struct References {
+    tables: Vec<String>,
+    functions: Vec<String>,
+    nested_with: bool,
+}
+
+impl Visitor for References {
+    type Break = ();
+
+    fn pre_visit_query(&mut self, query: &Query) -> ControlFlow<()> {
+        self.nested_with |= query.with.is_some();
+        ControlFlow::Continue(())
+    }
+
+    fn pre_visit_relation(&mut self, name: &ObjectName) -> ControlFlow<()> {
+        if name.0.len() == 1 {
+            if let Some(id) = name.0[0].as_ident() {
+                self.tables.push(if id.quote_style.is_some() {
+                    id.value.clone()
+                } else {
+                    id.value.to_lowercase()
+                });
+            }
+        }
+        ControlFlow::Continue(())
+    }
+
+    fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<()> {
+        if let Expr::Function(f) = expr {
+            self.functions.push(f.name.to_string().to_lowercase());
+        }
+        ControlFlow::Continue(())
+    }
+}
+
+type Part = (String, String, Vec<String>, Vec<String>);
+
+/// Return top-level CTEs and the final SELECT, with their direct references.
+#[pyfunction]
+#[pyo3(signature = (sql, allow_nested=false))]
+pub(crate) fn shared_ctes(sql: &str, allow_nested: bool) -> PyResult<Vec<Part>> {
+    let mut statements = Parser::parse_sql(&GenericDialect {}, sql)
+        .map_err(|e| PyValueError::new_err(e.to_string()))?;
+    if statements.len() != 1 {
+        return Err(PyValueError::new_err(
+            "shared execution requires one SELECT",
+        ));
+    }
+    let Statement::Query(mut query) = statements.remove(0) else {
+        return Err(PyValueError::new_err("shared execution requires a SELECT"));
+    };
+    let mut parts = Vec::new();
+    if let Some(with) = query.with.take() {
+        if with.recursive {
+            return Err(PyValueError::new_err("recursive CTEs cannot be staged"));
+        }
+        for cte in with.cte_tables {
+            if !cte.alias.columns.is_empty() {
+                return Err(PyValueError::new_err("CTE column aliases cannot be staged"));
+            }
+            let id = cte.alias.name;
+            let name = if id.quote_style.is_some() {
+                id.value
+            } else {
+                id.value.to_lowercase()
+            };
+            parts.push(part(name, &cte.query, allow_nested)?);
+        }
+    }
+    parts.push(part(String::new(), &query, allow_nested)?);
+    Ok(parts)
+}
+
+fn part(name: String, query: &Query, allow_nested: bool) -> PyResult<Part> {
+    let mut refs = References::default();
+    let _ = query.visit(&mut refs);
+    if refs.nested_with && !allow_nested {
+        return Err(PyValueError::new_err("nested WITH scopes cannot be staged"));
+    }
+    Ok((name, query.to_string(), refs.tables, refs.functions))
+}
+
+/// An internal Arrow stage with a verified ascending, NULL-last key order.
+/// The Python stage executor verifies or sorts these keys before registration
+/// and never mutates key buffers while a plan reads this provider.
+#[pyclass(module = "xarray_sql._native", name = "_SortedArrowTable")]
+pub(crate) struct SortedArrowTable {
+    table: Arc<dyn TableProvider>,
+}
+
+#[pymethods]
+impl SortedArrowTable {
+    #[new]
+    fn new(batch: &Bound<'_, PyAny>, keys: Vec<String>) -> PyResult<Self> {
+        let batch = RecordBatch::from_pyarrow_bound(batch)?;
+        let ordering = keys
+            .into_iter()
+            .map(|key| DFExpr::Column(Column::from_name(key)).sort(true, false))
+            .collect();
+        let table = MemTable::try_new(batch.schema(), vec![vec![batch]])
+            .map_err(|e| PyValueError::new_err(e.to_string()))?
+            .with_sort_order(vec![ordering]);
+        Ok(Self {
+            table: Arc::new(table),
+        })
+    }
+
+    fn __datafusion_table_provider__<'py>(
+        &self,
+        py: Python<'py>,
+        session: Bound<'py, PyAny>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let codec = crate::ffi_logical_codec_from_pycapsule(session)?;
+        let provider = FFI_TableProvider::new_with_ffi_codec(self.table.clone(), true, None, codec);
+        PyCapsule::new_with_value(py, provider, c"datafusion_table_provider")
+    }
+}
diff --git a/tests/test_ad.py b/tests/test_ad.py
index 0f868bc7..df3d2675 100644
--- a/tests/test_ad.py
+++ b/tests/test_ad.py
@@ -136,9 +136,54 @@ def test_a_mention_in_a_literal_is_not_a_call():
     assert rows == [{"s": "grad(x)", "jvp(v)": 1}]
 
 
-def test_prepare_refuses_a_gradient(problem):
-    ctx, *_ = problem
-    with pytest.raises(NotImplementedError, match="grad"):
-        ctx.prepare(
-            f"WITH {LOSS} SELECT feature, grad(loss, w) AS g FROM weights"
+def test_context_prepare_runs_gradient_steps_and_reads_writes(problem):
+    ctx, x, w, y = problem
+    step = ctx.prepare(
+        f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature"
+    )
+    old = step()
+    original = old["w"].to_numpy().copy()
+    replacement = xr.Dataset(
+        {"w": ("feature", w.copy() * 2)}, coords={"feature": np.arange(3)}
+    )
+    got = step(weights=replacement)
+    np.testing.assert_allclose(
+        got["w"], 2 * x.T @ (x @ replacement.w.values - y)
+    )
+    np.testing.assert_array_equal(old["w"], original)
+    with pytest.raises(ValueError, match="coordinates"):
+        step(weights=replacement.assign_coords(feature=[3, 4, 5]))
+
+
+def test_prepared_gradients_read_new_values_and_keep_old_results(problem):
+    ctx, x, w, y = problem
+    statements = [
+        f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature",
+        f"WITH {LOSS} SELECT feature, w - 0.01 * grad(loss, w) AS w FROM weights ORDER BY feature",
+    ]
+    run = ad.prepare(ctx, statements)
+    [first, _] = run()
+    original = first["w"].to_numpy().copy()
+    for scale in (0.5, -1.0, 2.0):
+        values = xr.Dataset(
+            {"w": ("feature", w.copy() * scale)},
+            coords={"feature": np.arange(3)},
         )
+        gradient, step = run(weights=values)
+        want = 2 * x.T @ (x @ values.w.values - y)
+        np.testing.assert_allclose(gradient["w"], want)
+        np.testing.assert_allclose(step["w"], values.w.values - 0.01 * want)
+        np.testing.assert_array_equal(first["w"], original)
+    # Regular ddx preparation resolves a replacement table by name.
+    ctx.deregister_table("weights")
+    ctx.from_dataset("weights", values * 2)
+    [gradient, _] = run()
+    np.testing.assert_allclose(
+        gradient["w"], 2 * x.T @ (x @ (values.w.values * 2) - y)
+    )
+
+
+def test_prepared_statements_require_a_sequence(problem):
+    ctx, *_ = problem
+    with pytest.raises(TypeError, match="list"):
+        ad.prepare(ctx, "SELECT 1")
diff --git a/tests/test_ml_demos.py b/tests/test_ml_demos.py
index b63ac452..06e72b21 100644
--- a/tests/test_ml_demos.py
+++ b/tests/test_ml_demos.py
@@ -1,7 +1,9 @@
 """The ML demos in ``benchmarks/ml`` run, at their smallest sizes.
 
-Each demo checks its results against JAX and exits with an error on a
-mismatch, so this runs each one and checks that it succeeds. It needs the
+The original demos check their results against JAX and exit on a
+mismatch. The Shakespeare trainer's end-to-end tests, using prepared
+temporary data, live in test_shakespeare.py; here its CLI is checked.
+The original demos need the
 ``ml`` extra, ddx's DuckDB backend, einfold and JAX, which the ``dev``
 dependency group installs.
 
@@ -24,6 +26,7 @@
         ["--steps", "1"],
         ["--steps", "1", "--optimizer", "adamw"],
     ],
+    "03_shakespeare.py": [["--help"]],
 }
 
 
diff --git a/tests/test_mlx.py b/tests/test_mlx.py
new file mode 100644
index 00000000..df02dfb9
--- /dev/null
+++ b/tests/test_mlx.py
@@ -0,0 +1,391 @@
+"""Public SQL and ddx contracts of the optional Apple silicon executor."""
+
+import os
+
+import numpy as np
+import pyarrow as pa
+import pytest
+import xarray as xr
+from datafusion import SessionConfig
+
+pytest.importorskip("mlx.core", exc_type=ImportError)
+if os.environ.get("MLX_ENABLE_TF32") != "0":
+    pytest.skip(
+        "GPU tests require MLX_ENABLE_TF32=0 at launch", allow_module_level=True
+    )
+pytest.importorskip("einfold")
+pytest.importorskip("ddxdb")
+
+from ddxdb import UnsupportedExpression
+
+from xarray_sql import Context, ad
+from xarray_sql.mlx import MLX
+
+PRODUCT = """
+SELECT a.i, b.k, SUM(a.v * b.v) AS v
+FROM a JOIN b ON a.j = b.j GROUP BY a.i, b.k ORDER BY a.i, b.k
+"""
+
+
+def _matrix(value, dims):
+    return xr.Dataset(
+        {"v": (dims, value)},
+        coords={d: np.arange(n) for d, n in zip(dims, value.shape)},
+        attrs={"experiment": "accelerated"},
+    )
+
+
+@pytest.fixture(params=["float64", "float32"])
+def accelerated(request):
+    try:
+        extension = MLX(
+            precision=request.param, min_ops=0, min_gpu_ops=0, query_block=3
+        )
+    except ImportError as exc:
+        pytest.skip(str(exc))
+    ctx = Context().with_extensions(extension)
+    return ctx, extension, 3e-5 if request.param == "float32" else 1e-12
+
+
+def _same(got, want, tolerance):
+    assert got.schema == want.schema
+    assert got.column_names == want.column_names
+    for g, w in zip(got.columns, want.columns):
+        if pa.types.is_floating(w.type):
+            assert g.null_count == w.null_count
+            assert g.is_null().to_pylist() == w.is_null().to_pylist()
+            np.testing.assert_allclose(
+                g.to_numpy(),
+                w.to_numpy(),
+                rtol=tolerance,
+                atol=tolerance,
+                equal_nan=True,
+            )
+            gv, wv = g.to_numpy(), w.to_numpy()
+            zero = (gv == 0) & (wv == 0)
+            np.testing.assert_array_equal(
+                np.signbit(gv[zero]), np.signbit(wv[zero])
+            )
+        else:
+            assert g.equals(w)
+
+
+def test_dense_shapes_row_order_and_current_values(accelerated):
+    ctx, extension, tolerance = accelerated
+    rng = np.random.default_rng(412)
+    for m, k, n in [(3, 7, 5), (17, 3, 1), (1, 9, 13), (23, 11, 19)]:
+        a, b = rng.normal(size=(m, k)), rng.normal(size=(k, n))
+        ctx.deregister_table("a")
+        ctx.deregister_table("b")
+        ctx.from_dataset("a", _matrix(a, ("i", "j")))
+        ctx.from_dataset("b", _matrix(b, ("j", "k")))
+        frame = ctx.sql(PRODUCT)
+        for scale in (1.0, -0.75):
+            a[:] *= scale
+            got = frame.to_dataset(dims=["i", "k"], template="a")
+            np.testing.assert_allclose(
+                got.v, a @ b, rtol=tolerance, atol=tolerance
+            )
+            assert got.attrs == {"experiment": "accelerated"}
+    device = "gpu" if extension.precision == "float32" else "cpu"
+    assert extension.statistics()[f"{device}_matmul"] > 0
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        "null",
+        "duplicate",
+        "missing",
+        "infinity",
+        "nan",
+        "empty",
+        "keys",
+        "signed_zero",
+        "extreme",
+    ],
+)
+def test_sql_bags_nulls_special_values_and_empty_inputs(accelerated, case):
+    ctx, _, tolerance = accelerated
+    a = {"i": [0, 0, 1, 1], "j": [0, 1, 0, 1], "v": [1.0, 2.0, 3.0, 4.0]}
+    b = {"j": [0, 0, 1, 1], "k": [0, 1, 0, 1], "v": [0.25, 0.5, 0.75, 1.0]}
+    if case == "null":
+        a["v"] = [None, None, 3.0, 4.0]
+    elif case == "duplicate":
+        for values in a.values():
+            values.append(values[0])
+    elif case == "missing":
+        for values in b.values():
+            values.pop()
+    elif case in ("infinity", "nan"):
+        a["v"][0] = np.inf if case == "infinity" else np.nan
+    elif case == "empty":
+        a = {key: [] for key in a}
+    elif case == "signed_zero":
+        a["v"] = [0.0, -0.0, -0.0, 0.0]
+    elif case == "extreme":
+        a["v"] = [1e300, -1e300, 1e-300, -1e-300]
+    elif case == "keys":
+        a["j"] = [None, -1, None, -1]
+        b["j"] = [None, None, -1, -1]
+    plain = Context()
+    for name, data in [("a", a), ("b", b)]:
+        batch = pa.record_batch(
+            data,
+            schema=pa.schema(
+                [
+                    (key, pa.float64() if key == "v" else pa.int64())
+                    for key in data
+                ]
+            ),
+        )
+        for con in (ctx, plain):
+            con.register_record_batches(name, [[batch]])
+    _same(
+        ctx.sql(PRODUCT).to_arrow_table(),
+        plain.sql(PRODUCT).to_arrow_table(),
+        tolerance,
+    )
+
+
+def test_prepared_forward_reads_writes_and_keeps_metadata(accelerated):
+    ctx, _, tolerance = accelerated
+    a = _matrix(np.ones((5, 3)), ("i", "j"))
+    b = _matrix(np.ones((3, 4)), ("j", "k"))
+    ctx.from_dataset("a", a)
+    ctx.from_dataset("b", b)
+    run = ctx.prepare(PRODUCT)
+    for scale in (1.0, 2.0, -0.5):
+        got = run(a=_matrix(np.full((5, 3), scale), ("i", "j")))
+        np.testing.assert_allclose(
+            got["v"], np.full(20, 3 * scale), atol=tolerance
+        )
+
+
+ATTENTION = """
+WITH scores AS (
+  SELECT q.i, k.j, SUM(q.v * k.v) / 2.0 AS s
+  FROM q JOIN k ON q.d = k.d WHERE {mask} GROUP BY q.i, k.j),
+weights AS (
+  SELECT i, j, exp(s - MAX(s) OVER (PARTITION BY i)) AS w FROM scores),
+norm AS (SELECT i, SUM(w) AS z FROM weights GROUP BY i)
+SELECT w.i, v.d, SUM(w.w / n.z * v.v) AS v
+FROM weights w JOIN v ON w.j = v.j JOIN norm n ON w.i = n.i
+GROUP BY w.i, v.d ORDER BY w.i, v.d
+"""
+
+
+@pytest.mark.parametrize(
+    "mask", ["q.i >= k.j", "q.i > k.j + 1", "q.i >= k.j AND q.i < k.j + 3"]
+)
+def test_attention_masks_are_absent_rows(accelerated, mask):
+    ctx, extension, tolerance = accelerated
+    plain = Context()
+    rng = np.random.default_rng(52)
+    datasets = [
+        (name, _matrix(rng.normal(size=(9, 4)), dims))
+        for name, dims in [
+            ("q", ("i", "d")),
+            ("k", ("j", "d")),
+            ("v", ("j", "d")),
+        ]
+    ]
+    for con in (ctx, plain):
+        for name, data in datasets:
+            con.from_dataset(name, data)
+    sql = ATTENTION.format(mask=mask)
+    _same(
+        ctx.sql(sql).to_arrow_table(),
+        plain.sql(sql).to_arrow_table(),
+        tolerance,
+    )
+    if extension.precision == "float32":
+        assert extension.statistics()["gpu_attention"] > 0
+    # A non-finite value in a masked row must not contaminate earlier rows.
+    datasets[-1][1]["v"].values[-1, 0] = np.inf
+    _same(
+        ctx.sql(sql).to_arrow_table(),
+        plain.sql(sql).to_arrow_table(),
+        tolerance,
+    )
+
+
+def test_ddx_grad_vjp_jvp_stop_gradient_and_repeated_steps(accelerated):
+    ctx, _extension, tolerance = accelerated
+    rng = np.random.default_rng(3)
+    a, b = rng.normal(size=(7, 4)), rng.normal(size=(4, 3))
+    ctx.from_dataset("a", _matrix(a, ("i", "j")))
+    ctx.from_dataset("b", _matrix(b, ("j", "k")))
+    loss = """WITH pred AS (
+      SELECT a.i, b.k, SUM(a.v * b.v) AS p FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k),
+      loss AS (SELECT SUM(p*p)/2 AS l FROM pred) """
+    select = "SELECT j,k,grad(loss,v) AS g FROM b ORDER BY j,k"
+    step = ctx.prepare(loss + select)
+    for scale in (1.0, -0.25, 2.0):
+        new = _matrix(b * scale, ("j", "k"))
+        rows = step(b=new)
+        np.testing.assert_allclose(
+            rows["g"].to_numpy().reshape(4, 3),
+            a.T @ (a @ new.v.values),
+            rtol=tolerance,
+            atol=tolerance,
+        )
+    # JVP uses ddx's tangent relation and the same accelerated contractions.
+    tangent = _matrix(np.ones((4, 3)), ("j", "k"))
+    ctx.from_dataset("tangent", tangent)
+    got = ctx.sql("""SELECT a.i,b.k,jvp(SUM(a.v*b.v),b.v,tangent) AS g
+      FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k ORDER BY i,k""").to_arrow_table()
+    np.testing.assert_allclose(
+        got["g"].to_numpy().reshape(7, 3),
+        a @ np.ones((4, 3)),
+        rtol=tolerance,
+        atol=tolerance,
+    )
+    ctx.sql(
+        "CREATE VIEW pred AS SELECT a.i,b.k,SUM(a.v*b.v) AS p FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k"
+    )
+    ctx.sql("CREATE VIEW cotangent AS SELECT i,k,1.0 AS p FROM pred")
+    got = ctx.sql(
+        "WITH pred AS (SELECT a.i,b.k,SUM(a.v*b.v) AS p FROM a "
+        "JOIN b ON a.j=b.j GROUP BY a.i,b.k) "
+        "SELECT j,k,vjp(pred,v,cotangent) AS g FROM b ORDER BY j,k"
+    ).to_arrow_table()
+    np.testing.assert_allclose(
+        got["g"].to_numpy().reshape(4, 3),
+        a.T @ np.ones((7, 3)),
+        rtol=tolerance,
+        atol=tolerance,
+    )
+    got = ctx.sql("""WITH loss AS (SELECT SUM(v*ddx_stop_gradient(v)) AS l FROM b)
+      SELECT j,k,grad(loss,v) AS g FROM b ORDER BY j,k""").to_arrow_table()
+    np.testing.assert_allclose(
+        got["g"].to_numpy().reshape(4, 3), b, rtol=tolerance, atol=tolerance
+    )
+
+
+def test_backend_threshold_retains_cpu_sql():
+    try:
+        extension = MLX(precision="float32", min_ops=10**12)
+    except ImportError as exc:
+        pytest.skip(str(exc))
+    ctx = Context(SessionConfig().with_target_partitions(2)).with_extensions(
+        extension
+    )
+    ctx.from_dataset("a", _matrix(np.ones((3, 4)), ("i", "j")))
+    ctx.from_dataset("b", _matrix(np.ones((4, 5)), ("j", "k")))
+    assert ctx.sql(PRODUCT).to_arrow_table()["v"].to_pylist() == [4.0] * 15
+    assert extension.statistics() == {}
+
+
+def test_batched_products_across_partitions(accelerated):
+    ctx, _, tolerance = accelerated
+    rng = np.random.default_rng(985)
+    a, b = rng.normal(size=(3, 5, 7)), rng.normal(size=(3, 7, 4))
+    ctx.from_dataset("a", _matrix(a, ("batch", "i", "j")), chunks={"batch": 1})
+    ctx.from_dataset("b", _matrix(b, ("batch", "j", "k")), chunks={"batch": 1})
+    got = ctx.sql("""SELECT a.batch,a.i,b.k,SUM(a.v*b.v) AS v
+      FROM a JOIN b ON a.batch=b.batch AND a.j=b.j
+      GROUP BY a.batch,a.i,b.k ORDER BY a.batch,a.i,b.k""").to_arrow_table()
+    np.testing.assert_allclose(
+        got["v"].to_numpy().reshape(3, 5, 4),
+        a @ b,
+        rtol=tolerance,
+        atol=tolerance,
+    )
+
+
+@pytest.mark.parametrize("mask", ["q.i >= k.j", "q.i < k.j + 2", "TRUE"])
+@pytest.mark.parametrize("per_head", [True, False])
+def test_attention_keeps_batches_and_heads_independent(
+    accelerated, mask, per_head
+):
+    ctx, extension, tolerance = accelerated
+    plain = Context()
+    rng = np.random.default_rng(542)
+    for name, shape, dims in [
+        ("q", (2, 3, 7, 4), ("seq", "h", "i", "d")),
+        ("k", (2, 3, 9, 4), ("seq", "h", "j", "d")),
+        ("v", (2, 3, 9, 5), ("seq", "h", "j", "e")),
+    ]:
+        data = _matrix(rng.normal(size=shape), dims)
+        ctx.from_dataset(name, data, chunks={"seq": 1})
+        plain.from_dataset(name, data, chunks={"seq": 1})
+    norm_keys = "seq,h,i" if per_head else "i"
+    norm_join = "a.seq=n.seq AND a.h=n.h AND a.i=n.i" if per_head else "a.i=n.i"
+    sql = f"""
+    WITH scores AS (
+      SELECT q.seq,q.h,q.i,k.j,SUM(q.v*k.v)/2 AS s FROM q JOIN k
+      ON q.seq=k.seq AND q.h=k.h AND q.d=k.d WHERE {mask}
+      GROUP BY q.seq,q.h,q.i,k.j),
+    weights AS (
+      SELECT seq,h,i,j,exp(s-MAX(s) OVER (PARTITION BY seq,h,i)) AS p
+      FROM scores),
+    norm AS (SELECT {norm_keys},SUM(p) AS z FROM weights GROUP BY {norm_keys})
+    SELECT a.seq,a.h,a.i,v.e,SUM(a.p/n.z*v.v) AS v FROM weights a
+    JOIN norm n ON {norm_join}
+    JOIN v ON a.seq=v.seq AND a.h=v.h AND a.j=v.j
+    GROUP BY a.seq,a.h,a.i,v.e ORDER BY a.seq,a.h,a.i,v.e
+    """
+    _same(
+        ctx.sql(sql).to_arrow_table(),
+        plain.sql(sql).to_arrow_table(),
+        tolerance,
+    )
+    if extension.precision == "float32" and per_head:
+        assert extension.statistics().get("gpu_attention", 0) > 0
+
+
+def test_ddx_unsupported_expression_still_raises(accelerated):
+    ctx, _, _ = accelerated
+    ctx.from_dataset(
+        "weights", xr.Dataset({"w": ("i", [1.2, 2.4])}, coords={"i": [0, 1]})
+    )
+    with pytest.raises(UnsupportedExpression, match="GROUP BY key"):
+        ctx.sql("""WITH buckets AS (
+          SELECT ROUND(w) AS k,SUM(w*w) AS s FROM weights GROUP BY ROUND(w)),
+          loss AS (SELECT SUM(k*s) AS l FROM buckets)
+          SELECT i,grad(loss,w) AS g FROM weights ORDER BY i""").collect()
+
+
+def test_attention_backward_uses_ddx_derivatives(accelerated):
+    ctx, _, tolerance = accelerated
+    rng = np.random.default_rng(123)
+    q, k, v = (rng.normal(size=(7, 4)) for _ in range(3))
+    for name, values, dims in [
+        ("q", q, ("i", "d")),
+        ("k", k, ("j", "d")),
+        ("v", v, ("j", "d")),
+    ]:
+        ctx.from_dataset(name, _matrix(values, dims))
+    scores = q @ k.T / 2
+    scores[np.triu_indices(7, 1)] = -np.inf
+    weights = np.exp(scores - np.max(scores, axis=1, keepdims=True))
+    weights /= weights.sum(axis=1, keepdims=True)
+    output = weights @ v
+    dw = output @ v.T
+    ds = weights * (dw - (dw * weights).sum(axis=1, keepdims=True))
+    want = [ds @ k / 2, ds.T @ q / 2, weights.T @ output]
+    model = ATTENTION.format(mask="q.i >= k.j")
+    statements = [
+        f"WITH attention AS ({model}), loss AS (SELECT SUM(v*v)/2 AS l FROM attention) "
+        f"SELECT {key},d,grad(loss,v) AS v FROM {name} ORDER BY {key},d"
+        for name, key in [("q", "i"), ("k", "j"), ("v", "j")]
+    ]
+    for frame, expected in zip(ad.sql_all(ctx, statements), want):
+        np.testing.assert_allclose(
+            frame.to_arrow_table()["v"].to_numpy().reshape(7, 4),
+            expected,
+            rtol=tolerance,
+            atol=tolerance,
+        )
+
+
+def test_low_backend_memory_limit_retains_cpu_sql():
+    extension = MLX(
+        precision="float32", min_ops=0, min_gpu_ops=0, max_working_bytes=1
+    )
+    ctx = Context().with_extensions(extension)
+    ctx.from_dataset("a", _matrix(np.ones((3, 4)), ("i", "j")))
+    ctx.from_dataset("b", _matrix(np.ones((4, 5)), ("j", "k")))
+    assert ctx.sql(PRODUCT).to_arrow_table()["v"].to_pylist() == [4.0] * 15
+    assert extension.statistics() == {}
diff --git a/tests/test_prepare.py b/tests/test_prepare.py
index 8cc356bc..1b8b2b53 100644
--- a/tests/test_prepare.py
+++ b/tests/test_prepare.py
@@ -9,6 +9,11 @@
 
 from xarray_sql import Context
 
+try:
+    from einfold.datafusion import EinFold
+except ImportError:
+    EinFold = None
+
 # The build side of the join, `a`, is what a reused hash join would keep.
 JOIN = """
 SELECT b.i, SUM(t.v2 * b.w) AS s
@@ -17,6 +22,26 @@
 """
 
 
+@pytest.mark.parametrize("extension", [False, True])
+def test_join_key_data_columns_can_change_their_range(extension):
+    ctx = Context()
+    if extension:
+        if EinFold is None:
+            pytest.skip("requires the optional EinFold planner")
+        ctx = ctx.with_extensions(EinFold())
+    a = xr.Dataset({"key": ("i", np.arange(128)), "v": ("i", np.ones(128))})
+    b = xr.Dataset({"key": ("j", np.arange(256)), "w": ("j", np.arange(256.0))})
+    ctx.from_dataset("a", a)
+    ctx.from_dataset("b", b)
+    sql = "SELECT a.i,b.j,a.v+b.w AS s FROM a JOIN b ON a.key=b.key ORDER BY a.i,b.j"
+    query = ctx.prepare(sql)
+    for start in (0, 128, 16):
+        updated = a.copy(deep=True)
+        updated["key"][:] = np.arange(start, start + 128)
+        rows = query(a=updated)
+        assert rows.equals(ctx.sql(sql).to_arrow_table())
+
+
 def _a(v):
     return xr.Dataset(
         {"v": (("i", "j"), v)},
@@ -262,3 +287,50 @@ def test_a_quoted_table_replaced_in_sql_counts_too():
     ctx.sql('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)').collect()
     with pytest.raises(ValueError, match="registered again"):
         query()
+
+
+def test_extension_handle_keeps_dataset_metadata_and_invalidation(ctx):
+    want = ctx.sql("SELECT * FROM a").to_dataset()
+    derived = ctx.with_extensions()
+    query = derived.prepare("SELECT i, j, v FROM a ORDER BY i, j")
+    xr.testing.assert_identical(
+        derived.sql("SELECT * FROM a").to_dataset(), want
+    )
+    ctx.deregister_table("a")
+    with pytest.raises(ValueError, match="registered again"):
+        query()
+
+
+@pytest.mark.parametrize(
+    "chunks", [{"i": 1, "j": 1}, {"i": 1}, {"i": 4, "j": 5}]
+)
+def test_fine_chunk_registration_keeps_rows_metadata_and_writes(chunks):
+    values = np.arange(129 * 33, dtype=np.float64).reshape(129, 33)
+    ds = _a(values)
+    ds.attrs["experiment"] = "fine chunks"
+    ctx = Context()
+    ctx.from_dataset("a", ds, chunks=chunks)
+    sql = "SELECT i,j,v FROM a ORDER BY i,j"
+    xr.testing.assert_identical(ctx.sql(sql).to_dataset(), ds)
+    run = ctx.prepare(sql)
+    replacement = ds * 2
+    rows = run(a=replacement)
+    np.testing.assert_array_equal(rows["v"], replacement.v.values.reshape(-1))
+
+
+@pytest.mark.parametrize("threads", [1, 60])
+def test_thread_target_is_explicit_and_keeps_xarray_support(threads):
+    ctx = Context(threads=threads)
+    assert ctx.sql(
+        "SELECT value FROM information_schema.df_settings "
+        "WHERE name='datafusion.execution.target_partitions'"
+    ).to_pylist() == [{"value": str(threads)}]
+    ctx.from_dataset("a", _a(np.ones((4, 3))))
+    assert ctx.prepare("SELECT SUM(v) AS s FROM a")().to_pylist() == [
+        {"s": 12.0}
+    ]
+
+
+def test_thread_target_must_be_positive():
+    with pytest.raises(ValueError, match="positive"):
+        Context(threads=0)
diff --git a/tests/test_shared.py b/tests/test_shared.py
new file mode 100644
index 00000000..07a36905
--- /dev/null
+++ b/tests/test_shared.py
@@ -0,0 +1,296 @@
+"""Shared SELECT stages retain SQL row semantics and original ddx lineage."""
+
+import numpy as np
+import pyarrow as pa
+import pytest
+import xarray as xr
+
+from xarray_sql import Context
+
+pytest.importorskip("ddxdb")
+einfold = pytest.importorskip("einfold._native")
+pytestmark = pytest.mark.skipif(
+    not getattr(einfold, "RESET_ON_DECODE", False),
+    reason="requires the optional EinFold prepared-plan patch",
+)
+
+from einfold.datafusion import EinFold
+
+from xarray_sql.shared import SharedExecutor
+
+try:
+    from xarray_sql.dense import DenseExecutor
+except ImportError:
+    DenseExecutor = None
+
+_EXECUTORS = [
+    SharedExecutor,
+    pytest.param(
+        DenseExecutor,
+        marks=pytest.mark.skipif(
+            DenseExecutor is None, reason="requires compact EinFold programs"
+        ),
+    ),
+]
+
+
+@pytest.fixture(params=_EXECUTORS)
+def executor(request):
+    ctx = Context().with_extensions(EinFold())
+    ctx.from_dataset("w", xr.Dataset({"v": ("i", [1.0, 2.0, 3.0])}))
+    return request.param(ctx)
+
+
+def test_shared_join_refreshes_build_side_and_returns_snapshots(executor):
+    sql = """
+    WITH square AS (SELECT i, v*v AS s FROM w),
+    doubled AS (SELECT a.i, a.s + b.s AS s FROM square a
+                JOIN square b ON a.i = b.i)
+    SELECT * FROM doubled ORDER BY i
+    """
+    query = executor.prepare(sql)
+    first = query()
+    assert first.to_pydict() == {"i": [0, 1, 2], "s": [2.0, 8.0, 18.0]}
+    for values in ([3.0, 2.0, 1.0], [5.0, 6.0, 7.0]):
+        rows = query(w=xr.Dataset({"v": ("i", values)}))
+        assert rows.equals(executor.ctx.sql(sql).to_arrow_table())
+    assert first.to_pydict()["s"] == [2.0, 8.0, 18.0]
+
+
+def test_virtual_view_gradients_use_original_dependencies(executor):
+    executor.register_view("model", "SELECT i, v*v AS s FROM w")
+    loss = "WITH loss AS (SELECT SUM(s) AS l FROM model)"
+    forward = executor.prepare("SELECT SUM(s) AS l FROM model")
+    step = executor.prepare_all(
+        [
+            loss + " SELECT i, grad(loss, v) AS g FROM w ORDER BY i",
+            loss + " SELECT i, v - 0.1*grad(loss, v) AS v FROM w ORDER BY i",
+        ]
+    )
+    for values in ([1.0, 2.0, 3.0], [4.0, -2.0, 0.0]):
+        data = xr.Dataset({"v": ("i", values)})
+        assert forward(w=data)["l"].to_pylist() == [np.dot(values, values)]
+        gradient, updated = step()
+        np.testing.assert_allclose(gradient["g"], 2 * np.array(values))
+        np.testing.assert_allclose(updated["v"], 0.8 * np.array(values))
+
+
+def test_gradients_through_reused_ctes_in_a_virtual_view(executor):
+    executor.register_view(
+        "model",
+        """
+        WITH squared AS (SELECT i, v*v AS a FROM w),
+        mean AS (SELECT AVG(a) AS m FROM squared),
+        centered AS (SELECT i, power(a-m, 2) AS s FROM squared CROSS JOIN mean),
+        total AS (SELECT SUM(s) AS t FROM centered)
+        SELECT i, s+t AS s FROM centered CROSS JOIN total
+    """,
+    )
+    step = executor.prepare_all(
+        [
+            (
+                "WITH loss AS (SELECT SUM(s) AS l FROM model) "
+                "SELECT i, grad(loss,v) AS g FROM w ORDER BY i"
+            )
+        ]
+    )
+    for values in ([1.0, 2.0, 3.0], [-2.0, 4.0, 1.0]):
+        v = np.array(values)
+        [rows] = step(w=xr.Dataset({"v": ("i", v)}))
+        np.testing.assert_allclose(rows["g"], 16 * v * (v * v - np.mean(v * v)))
+
+
+def test_registration_and_coordinate_changes_refresh_plans(executor):
+    query = executor.prepare(
+        "WITH t AS (SELECT i, v FROM w WHERE v > 0) "
+        "SELECT a.i, a.v+b.v AS v FROM t a JOIN t b ON a.i=b.i ORDER BY a.i"
+    )
+    assert query().num_rows == 3
+    for coords, values in (([0, 1, 2], [-1.0, 2.0, 4.0]), ([8, 9], [5.0, 6.0])):
+        executor.ctx.deregister_table("w")
+        executor.ctx.from_dataset(
+            "w", xr.Dataset({"v": ("i", values)}, coords={"i": coords})
+        )
+        assert query().to_pydict() == {
+            "i": [i for i, v in zip(coords, values) if v > 0],
+            "v": [2 * v for v in values if v > 0],
+        }
+
+
+def test_empty_stages_can_gain_and_lose_rows_after_inplace_writes(executor):
+    query = executor.prepare(
+        "WITH t AS (SELECT i,v FROM w WHERE v < 0) "
+        "SELECT a.i,a.v+b.v AS v FROM t a JOIN t b ON a.i=b.i ORDER BY a.i"
+    )
+    for values in (
+        [1.0, 2.0, 3.0],
+        [2.0, 3.0, 4.0],
+        [-1.0, -2.0, 3.0],
+        [1.0, 2.0, 3.0],
+    ):
+        rows = query(w=xr.Dataset({"v": ("i", values)}))
+        assert rows.to_pydict() == {
+            "i": [i for i, v in enumerate(values) if v < 0],
+            "v": [2 * v for v in values if v < 0],
+        }
+
+
+def test_changed_domains_refresh_every_shared_stage(executor):
+    query = executor.prepare("""
+        WITH a AS (SELECT i,v FROM w WHERE v < 0),
+        b AS (SELECT x.i,x.v+y.v AS v FROM a x JOIN a y ON x.i=y.i),
+        c AS (SELECT x.i,x.v+y.v AS v FROM b x JOIN b y ON x.i=y.i)
+        SELECT x.i,x.v+y.v AS v FROM c x JOIN c y ON x.i=y.i ORDER BY x.i
+    """)
+    for values in (
+        [1.0, 2.0, 3.0],
+        [-1.0, -2.0, 3.0],
+        [1.0, -3.0, -4.0],
+        [1.0, 2.0, 3.0],
+    ):
+        rows = query(w=xr.Dataset({"v": ("i", values)}))
+        assert rows.to_pydict() == {
+            "i": [i for i, v in enumerate(values) if v < 0],
+            "v": [8 * v for v in values if v < 0],
+        }
+
+
+@pytest.mark.parametrize("executor_type", _EXECUTORS)
+def test_shared_stages_keep_nulls_and_duplicate_rows(executor_type):
+    ctx = Context().with_extensions(EinFold())
+    ctx.register_record_batches(
+        "w",
+        [
+            [
+                pa.record_batch(
+                    {
+                        "i": [0, 0, 1],
+                        "v": [1.0, None, 3.0],
+                    }
+                )
+            ]
+        ],
+    )
+    executor = executor_type(ctx)
+    sql = (
+        "WITH t AS (SELECT i, v FROM w) "
+        "SELECT a.i, a.v AS a, b.v AS b FROM t a JOIN t b ON a.i=b.i "
+        "ORDER BY a.i, a.v NULLS FIRST, b.v NULLS FIRST"
+    )
+    query = executor.prepare(sql)
+    for _ in range(2):
+        assert query().equals(ctx.sql(sql).to_arrow_table())
+
+
+def test_optimizer_scalar_functions_remain_available(executor):
+    rows = executor.prepare(
+        "WITH norm AS (SELECT sqrt(SUM(v*v)) AS n FROM w), "
+        "clip AS (SELECT least(1.0, 1.0/(n+1e-6)) AS c FROM norm) "
+        "SELECT i, v*clip.c AS v FROM w CROSS JOIN clip ORDER BY i"
+    )()
+    np.testing.assert_allclose(
+        rows["v"], np.array([1.0, 2.0, 3.0]) / (np.sqrt(14) + 1e-6)
+    )
+
+
+def test_ordered_stages_preserve_quoted_names_and_window_results(executor):
+    sql = """
+    WITH t AS (SELECT i AS "axis.name", v FROM w),
+    doubled AS (SELECT a."axis.name", a.v+b.v AS v FROM t a JOIN t b
+                ON a."axis.name"=b."axis.name")
+    SELECT "axis.name", SUM(v) OVER (
+      ORDER BY "axis.name" ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS s
+    FROM doubled ORDER BY "axis.name"
+    """
+    query = executor.prepare(sql)
+    for values in ([3.0, 1.0, 2.0], [5.0, -1.0, 4.0]):
+        rows = query(w=xr.Dataset({"v": ("i", values)}))
+        assert rows.equals(executor.ctx.sql(sql).to_arrow_table())
+
+
+def test_many_gradient_contributions_share_a_bounded_reduction_graph(executor):
+    terms = " UNION ALL ".join(
+        f"SELECT i,power(v,{p}) AS y FROM w" for p in range(2, 42)
+    )
+    step = executor.prepare_all(
+        [
+            (
+                f"WITH terms AS ({terms}), loss AS (SELECT SUM(y) AS l FROM terms) "
+                "SELECT i,grad(loss,v) AS g FROM w ORDER BY i"
+            )
+        ]
+    )
+    for values in ([1.0, 2.0, 3.0], [2.0, -1.0, 0.0]):
+        v = np.array(values)
+        [rows] = step(w=xr.Dataset({"v": ("i", v)}))
+        np.testing.assert_allclose(
+            rows["g"], sum(p * v ** (p - 1) for p in range(2, 42))
+        )
+
+
+@pytest.mark.parametrize("executor_type", _EXECUTORS)
+def test_matrix_gradients_through_a_reconstructed_normalized_operand(
+    executor_type,
+):
+    rng = np.random.default_rng(48)
+    a, b, g = (
+        rng.normal(size=(4, 3)),
+        rng.normal(size=(3, 5)),
+        rng.normal(size=3),
+    )
+    ctx = Context().with_extensions(EinFold())
+    ctx.from_dataset("a", xr.Dataset({"v": (("i", "j"), a)}))
+    ctx.from_dataset("b", xr.Dataset({"w": (("j", "k"), b)}))
+    ctx.from_dataset("g", xr.Dataset({"gain": ("j", g)}))
+    executor = executor_type(ctx)
+    executor.register_view(
+        "normalized",
+        """
+        WITH mean AS (SELECT i,AVG(v) AS m FROM a GROUP BY i),
+        variance AS (SELECT a.i,AVG(power(a.v-mean.m,2)) AS v
+                     FROM a JOIN mean ON a.i=mean.i GROUP BY a.i)
+        SELECT a.i,a.j,(a.v-mean.m)/sqrt(variance.v+1e-5)*g.gain AS v
+        FROM a JOIN mean ON a.i=mean.i JOIN variance ON a.i=variance.i
+        JOIN g ON a.j=g.j
+    """,
+    )
+    step = executor.prepare_all(
+        [
+            """
+        WITH predicted AS (SELECT x.i,b.k,SUM(x.v*b.w) AS v
+              FROM normalized x JOIN b ON x.j=b.j GROUP BY x.i,b.k),
+        loss AS (SELECT SUM(v*v)/2 AS l FROM predicted)
+        SELECT j,k,grad(loss,w) AS g FROM b ORDER BY j,k
+    """
+        ]
+    )
+    centered = a - a.mean(axis=1, keepdims=True)
+    normalized = (
+        centered
+        / np.sqrt(np.mean(centered**2, axis=1, keepdims=True) + 1e-5)
+        * g
+    )
+    for scale in (1.0, -0.25):
+        current = b * scale
+        expected = normalized.T @ (normalized @ current)
+        [rows] = step(b=xr.Dataset({"w": (("j", "k"), current)}))
+        np.testing.assert_allclose(
+            rows["g"].to_numpy().reshape(3, 5),
+            expected,
+            rtol=1e-12,
+            atol=1e-12,
+        )
+
+
+@pytest.mark.parametrize(
+    "sql",
+    [
+        "WITH t AS (SELECT random() AS x) SELECT * FROM t",
+        "WITH RECURSIVE t AS (SELECT 1) SELECT * FROM t",
+        "WITH t AS (WITH s AS (SELECT 1) SELECT * FROM s) SELECT * FROM t",
+        "SELECT 1; SELECT 2",
+    ],
+)
+def test_unsupported_staging_fails_explicitly(executor, sql):
+    with pytest.raises(ValueError):
+        executor.prepare(sql)
diff --git a/uv.lock b/uv.lock
index 50583113..508b9263 100644
--- a/uv.lock
+++ b/uv.lock
@@ -622,14 +622,14 @@ wheels = [
 
 [[package]]
 name = "einfold"
-version = "0.1.0rc2"
+version = "0.2.0"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/76/e9/b534c6660e42ae62f3bd430611b299f406833cd82757e39f74d0161290e6/einfold-0.1.0rc2.tar.gz", hash = "sha256:82df6a9e2e1954cc3ba90982b61a1236b0451e84a5f80ab4e7707c0e46f7950e", size = 224290, upload-time = "2026-10-30T02:08:18.251Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/3d/03/7ea71305a4aa642c26224e598c49923a1f074659d479589c744c304a3eea/einfold-0.2.0.tar.gz", hash = "sha256:fe4e6970e8e04a53b0f4d180375e160a27f015b1f4f38294d94338468fa159c5", size = 266243, upload-time = "2026-11-01T01:35:47.288Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/9a/19/986bd602078909cfa9bba3133d6f9cf5372c0e9f49de891aeeb657334187/einfold-0.1.0rc2-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:907bd43840bd3f0006436705a1f9a62939c2e8410e10911753c72975ce8edba0", size = 43853038, upload-time = "2026-10-30T02:07:57.966Z" },
-    { url = "https://files.pythonhosted.org/packages/46/56/3045c7efff5a649f8f209d95175785d7ddb7a279925f2033627a62df0645/einfold-0.1.0rc2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:1041aebdd7293dac48b680addfebb6bd90870e220a56263fe7fb43da4f3b1217", size = 42234075, upload-time = "2026-10-30T02:08:00.816Z" },
-    { url = "https://files.pythonhosted.org/packages/d2/bf/9d73c11836034eceea3057e246e89beadc0b090b11c4bb88865e357102af/einfold-0.1.0rc2-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:55d10e500ef934322e1108fa24b85da37ce21b152a9b3dd0364bf604e93fb0e0", size = 48676482, upload-time = "2026-10-30T02:08:03.849Z" },
-    { url = "https://files.pythonhosted.org/packages/86/63/8dc06cb2ffc0e8bc982b83d1ffd60b9bba30f7a5990c0f11254f9faa5825/einfold-0.1.0rc2-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:98cc8014c4e1f26101ecd492e9881ff3bafdddcf130692fa97a600a3f8047d6d", size = 47269964, upload-time = "2026-10-30T02:08:06.821Z" },
+    { url = "https://files.pythonhosted.org/packages/97/a9/e3b22d401da30443d6aa42b4bcf21fce4736669a6a0b59368bc07f59b4fa/einfold-0.2.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a1c4aa1f4177c0079088bfe2c4c9fc50fb7a5f34af7cbc8251bd4376d8740014", size = 43892338, upload-time = "2026-11-01T01:35:30.093Z" },
+    { url = "https://files.pythonhosted.org/packages/03/f3/8e3125e25165eee2042f7da09806a980b1703d9daec51a0fdfcf22eba9a1/einfold-0.2.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:e0572d2282bb283012b814429ade84f395c8c72732ee2989c9d40093e9bf3031", size = 42425356, upload-time = "2026-11-01T01:35:34.403Z" },
+    { url = "https://files.pythonhosted.org/packages/c5/69/a19afa99bcc6dfab3ef57411cda214e3aedd1eae8a8cfa7fd0fa1f7df754/einfold-0.2.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0d2dedfdbbc6f67f8864554e73c9a9718d85ae27c7e0cfc8ce547afddbb64239", size = 48717033, upload-time = "2026-11-01T01:35:39.06Z" },
+    { url = "https://files.pythonhosted.org/packages/31/cf/86eb4757245705709f78d66982d90a6bff47ebf294d0777e4f6e8c3810ac/einfold-0.2.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:c2458223021307aaf3bff213954a1b7cad301eab51af16f6e24774dc0505a2d8", size = 47308031, upload-time = "2026-11-01T01:35:44.234Z" },
 ]
 
 [package.optional-dependencies]
@@ -1573,6 +1573,53 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/5a/ff/bda40387b5c5c64254595f4d81a12351770856acc5de4e6d43606a31f161/ml_dtypes-0.6.0-cp315-cp315t-win_arm64.whl", hash = "sha256:f6cb525101b6b903779188c1e9e9490c343b455ab822883e02cf01e5547338d2", size = 572666, upload-time = "2026-08-13T14:14:38.993Z" },
 ]
 
+[[package]]
+name = "mlx"
+version = "0.32.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mlx-metal" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/52/73/f2b18f4bbcfeacd6880e68f2854609d6d2d305e5676586451a6c26664061/mlx-0.32.4-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:60ba41564e8b5018c9d8a5f8ce7a5aaf70c96d8dd98c1a54b854cbada362d471", size = 671799, upload-time = "2026-10-14T21:15:37.06Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/f0/5a722aa39986cf274438f0acc22e90304cdafdb7265045cd893924a5608f/mlx-0.32.4-cp310-cp310-macosx_15_0_arm64.whl", hash = "sha256:e635385a4f778b248ea6f61fbeb285b6bb14a1891d75cfbfd5af12e010298308", size = 671801, upload-time = "2026-10-14T21:15:39.02Z" },
+    { url = "https://files.pythonhosted.org/packages/95/27/9069458b0079820b860678f28fcb6bd0a3998e17db35440fc9b2847d2145/mlx-0.32.4-cp310-cp310-macosx_26_0_arm64.whl", hash = "sha256:6dd2d7948c3708a62a1b13c32b03b5a94031e0cd3b7f5c5a0343e5a4a3b3bc41", size = 671791, upload-time = "2026-10-14T21:15:40.697Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/fd/f7af63951d6f4d39b47bbda462f2ed0341ce1293967de36ce9bbd4b836a9/mlx-0.32.4-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:dd0c9070d9c362fe280a0224806b1b47719b154984735041ebcc907e597c08ca", size = 671764, upload-time = "2026-10-14T21:15:48.321Z" },
+    { url = "https://files.pythonhosted.org/packages/56/1d/e23a150073e07212c23358b978af0bfd21b87f0b97937ec1f300cee346ce/mlx-0.32.4-cp311-cp311-macosx_15_0_arm64.whl", hash = "sha256:fbb3fcf1363a631d05aa5adcb30d07655158be07e09339c1fe62f7267b3a78f5", size = 671761, upload-time = "2026-10-14T21:15:49.94Z" },
+    { url = "https://files.pythonhosted.org/packages/c7/9b/3e5cfd79fd468dbc7f516e5fa7341419541804c5f8beefe04053326e25e9/mlx-0.32.4-cp311-cp311-macosx_26_0_arm64.whl", hash = "sha256:dfd5269398bb64372137d8c354530b6e81159ce0e92fd1ed7e9140b294fd939a", size = 671758, upload-time = "2026-10-14T21:15:51.859Z" },
+    { url = "https://files.pythonhosted.org/packages/69/b0/e726ce51e1b445e4ed08c2d0e036a01c7532f4d1a856328da546a6f883bd/mlx-0.32.4-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:ff1d7be54758ea3a875b8fe54eff13bed17ad9566505e301f5d7851bf869680f", size = 667597, upload-time = "2026-10-14T21:16:00.688Z" },
+    { url = "https://files.pythonhosted.org/packages/6d/92/7a999dcf6460525a31e555905cdca2f835ec0d786c528e969e04cd00d158/mlx-0.32.4-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:f0604cb3015091f285c79c7b7dbfcabcfa9948248728c7bab6e0263513c3d3fe", size = 667599, upload-time = "2026-10-14T21:16:02.467Z" },
+    { url = "https://files.pythonhosted.org/packages/af/44/90bbdc3cdd120e1181f7c4f617ff996530159f2ea3ba3c08b83b08314d21/mlx-0.32.4-cp312-cp312-macosx_26_0_arm64.whl", hash = "sha256:7928869076db6267ab3ebdb253b2c80be1e738b3032c783170bd1d59d460c578", size = 667622, upload-time = "2026-10-14T21:16:04.219Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/25/b570700cb3d7b37477a865382880a6549b48b9efd2dd356dbc16eaaa01d1/mlx-0.32.4-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:c7375f2d6acd8c72f413d018266ae06387cb9caab31f475f3ff1c5600a7cb4c5", size = 667579, upload-time = "2026-10-14T21:16:12.885Z" },
+    { url = "https://files.pythonhosted.org/packages/1c/ad/1377a281f1b81f42929a3887b5225d0a38fba2d804ac1039011a223d2ab2/mlx-0.32.4-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:d46dd07f57a65bae00ba7da239ee76ebf87e3202f042808d209b56be3140e416", size = 667585, upload-time = "2026-10-14T21:16:14.597Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/b1/33ca08f84550fd512f28b84055ff00631ef2ce83011ac2bbc171ffec7bbb/mlx-0.32.4-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:ef130e3da895c85374b11a6c389acb122893f2ffcbed826b8b62fd45a20b7ad6", size = 667575, upload-time = "2026-10-14T21:16:16.31Z" },
+    { url = "https://files.pythonhosted.org/packages/98/85/f020b6392f04e37e9d2263daa194061736801ad51af06b1f11a253643976/mlx-0.32.4-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:43a6469797f09a985f8f035df63cd072105d5f16bfbe27d82b542065c988c3dd", size = 692145, upload-time = "2026-10-14T21:16:24.976Z" },
+    { url = "https://files.pythonhosted.org/packages/df/db/95168f9492b284042628952778bfc2342c0f381ac2434b3640b80fa2f292/mlx-0.32.4-cp313-cp313t-macosx_15_0_arm64.whl", hash = "sha256:986f001874378d0a4c1f919817b7e331be96bbebf3993f5182459bcfc6ac0116", size = 692143, upload-time = "2026-10-14T21:16:26.705Z" },
+    { url = "https://files.pythonhosted.org/packages/fc/81/c452f266d36c5708182ce079c2a6b6beca5c3405279216b28510368d80a7/mlx-0.32.4-cp313-cp313t-macosx_26_0_arm64.whl", hash = "sha256:733d28fc6b200910211924cd8d5deef86095cd085a0bf64502943c80b22942d5", size = 692095, upload-time = "2026-10-14T21:16:28.371Z" },
+    { url = "https://files.pythonhosted.org/packages/99/2f/ed874a2de4ad30b47f27960373880aa24b53d8963ea81125e12e460340c2/mlx-0.32.4-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f1e6ef441e1e2ce54159022a8fa5597f8ea4f7acf278752049c4f774bbd0501d", size = 668901, upload-time = "2026-10-14T21:16:36.971Z" },
+    { url = "https://files.pythonhosted.org/packages/12/c6/1e15326918233e1ee74f55bcc5577f6cce1c4257e5a1bbadcdb30b589dca/mlx-0.32.4-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:83adc889ae89057fd9f292ca1a9ee56da43f181fccc7ae48a060d7aadc179eaa", size = 668899, upload-time = "2026-10-14T21:16:39.057Z" },
+    { url = "https://files.pythonhosted.org/packages/95/2f/f84f5f16000adbdb1c53583f53da045165013ebb62f6943c169714be011f/mlx-0.32.4-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:1dd56491b44a2bfa217f60da66a5c9c889add7ad9dc2cbcf526389ea31d83bee", size = 668817, upload-time = "2026-10-14T21:16:40.738Z" },
+    { url = "https://files.pythonhosted.org/packages/e3/7c/c90d57df7775421fe43f7a235f6cee85a67c4aecc7325051b23fff58335b/mlx-0.32.4-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f6cbda9cb8c33c9f8673a2b69fb19bdab75b5032d4e47a42a0e7b36db1c3ebcc", size = 692106, upload-time = "2026-10-14T21:16:50.019Z" },
+    { url = "https://files.pythonhosted.org/packages/6e/07/90e629a9667ec35b2aa8c84e0e096ff36fdd58731dc9d182f5ee59a488d0/mlx-0.32.4-cp314-cp314t-macosx_15_0_arm64.whl", hash = "sha256:4f95afac25431585b34846d472ea0ad5db83fc97fb014d62c70082456f6b6b5c", size = 692107, upload-time = "2026-10-14T21:16:51.978Z" },
+    { url = "https://files.pythonhosted.org/packages/d0/20/bdb26b961fc8f0d54e6b36dd5360449f0cf8591d14cee10470c14ce1e154/mlx-0.32.4-cp314-cp314t-macosx_26_0_arm64.whl", hash = "sha256:63ac6fde8a55d688448fc25eb32b645276501e0a30ff961a72958dfd0a4065df", size = 692103, upload-time = "2026-10-14T21:16:53.963Z" },
+    { url = "https://files.pythonhosted.org/packages/74/d7/6abf128b1740d8be1d711d8ded6ccd4714c53d8af8b6267ac6aae285c60e/mlx-0.32.4-cp315-cp315-macosx_14_0_arm64.whl", hash = "sha256:4f8e640fd3afcc19da343c5a769743f361573b8db98cd739e295a387921486a5", size = 669191, upload-time = "2026-10-14T21:17:02.83Z" },
+    { url = "https://files.pythonhosted.org/packages/43/b7/3b9bcbb9a5d32025d032fb933b7c1a92bf1a4e22a796ff3b3ef7c2a7a53c/mlx-0.32.4-cp315-cp315-macosx_15_0_arm64.whl", hash = "sha256:6c6cf406fe093b4ea21876949e0d21f75384900f3df543e8411447ee37cdd394", size = 669189, upload-time = "2026-10-14T21:17:04.66Z" },
+    { url = "https://files.pythonhosted.org/packages/45/10/1164007669c9e015e1124d1387ba03079ffb6650224cbc01aaba1a43ebfd/mlx-0.32.4-cp315-cp315-macosx_26_0_arm64.whl", hash = "sha256:869c4e2d99ef622cfb39e6581afe5d3286e0c01161e03ad580561db28bcd3992", size = 669097, upload-time = "2026-10-14T21:17:06.647Z" },
+    { url = "https://files.pythonhosted.org/packages/a3/21/8b1b2f90bb3fa27d4bb5d0bbfb07c7f210844da6bb89f88476ab0ba84f50/mlx-0.32.4-cp315-cp315t-macosx_14_0_arm64.whl", hash = "sha256:95830850ababacf86d4e1e11b8cfa062a14b12c0c5f332dad1d4f99e1679a162", size = 692273, upload-time = "2026-10-14T21:17:16.639Z" },
+    { url = "https://files.pythonhosted.org/packages/a5/90/a5f916a6937481cc5ebab3f461d4a3459175d0ddf7e0339ca15b97a42fc8/mlx-0.32.4-cp315-cp315t-macosx_15_0_arm64.whl", hash = "sha256:09da967a668de8b7033de3d74930efcf36ef993ede93173f51f11452a1c09b22", size = 692274, upload-time = "2026-10-14T21:17:18.385Z" },
+    { url = "https://files.pythonhosted.org/packages/47/3e/7371d6df935a531dfcd1307c28d11d8308f8304bec25d0575a36e6c451e7/mlx-0.32.4-cp315-cp315t-macosx_26_0_arm64.whl", hash = "sha256:1e52124110303abb7513e6575f59c4c1d32e721c1a45f7a2ae200e8cc492ef13", size = 692240, upload-time = "2026-10-14T21:17:20.391Z" },
+]
+
+[[package]]
+name = "mlx-metal"
+version = "0.32.4"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/77/24/7e5b7ed5d28abacab4577c7a8dd1fca6e9058504f5f95ae4333c09420e2e/mlx_metal-0.32.4-py3-none-macosx_14_0_arm64.whl", hash = "sha256:8019b2db622599267d66cc168b2f0dbf5af929790a6c3b3913e00eabc42a7947", size = 44958619, upload-time = "2026-10-14T21:14:45.727Z" },
+    { url = "https://files.pythonhosted.org/packages/54/8c/9754dd6c4db485037f59cc8be6495738e14cdeb8bf19d5bddb5e96da890f/mlx_metal-0.32.4-py3-none-macosx_15_0_arm64.whl", hash = "sha256:9f0a3cc26d89724434d9684932fc6f3e4e83669d2a895d858ed2c8a55eb84e8b", size = 44942974, upload-time = "2026-10-14T21:14:50.012Z" },
+    { url = "https://files.pythonhosted.org/packages/2a/d6/4340ed6e184f71693d7f531c9ae55d677451351245760d71f8dd714dcbe5/mlx_metal-0.32.4-py3-none-macosx_26_0_arm64.whl", hash = "sha256:514df92d41a286895d2564b4b18f4da45096bd7a4a0c47e1042f345ad5a7f936", size = 78106861, upload-time = "2026-10-14T21:14:56.184Z" },
+]
+
 [[package]]
 name = "multidict"
 version = "6.6.3"
@@ -3132,6 +3179,10 @@ geo = [
 ml = [
     { name = "ddxdb", marker = "sys_platform != 'emscripten'" },
 ]
+mlx = [
+    { name = "einfold", extra = ["datafusion"], marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" },
+    { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" },
+]
 polars = [
     { name = "polars" },
 ]
@@ -3175,8 +3226,10 @@ requires-dist = [
     { name = "datafusion", marker = "sys_platform != 'emscripten'", specifier = "==55.0.0" },
     { name = "ddxdb", marker = "sys_platform != 'emscripten' and extra == 'ml'", specifier = ">=0.4.6,<0.5" },
     { name = "duckdb", marker = "extra == 'duckdb'", specifier = ">=1.4.0" },
+    { name = "einfold", extras = ["datafusion"], marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'mlx'", specifier = ">=0.2,<0.3" },
     { name = "gcsfs", marker = "extra == 'test'" },
     { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'" },
+    { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'mlx'", specifier = ">=0.32.4,<0.33" },
     { name = "polars", marker = "extra == 'polars'", specifier = ">=1.33" },
     { name = "pre-commit", marker = "extra == 'dev'" },
     { name = "pyarrow", specifier = ">=16.0.0" },
@@ -3190,12 +3243,12 @@ requires-dist = [
     { name = "xarray-sql", extras = ["docs"], marker = "extra == 'dev'" },
     { name = "zensical", marker = "extra == 'docs'" },
 ]
-provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "ml", "polars", "test"]
+provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "ml", "mlx", "polars", "test"]
 
 [package.metadata.requires-dev]
 dev = [
     { name = "ddxdb", extras = ["duckdb"], marker = "sys_platform != 'emscripten'", specifier = ">=0.4.6,<0.5" },
-    { name = "einfold", extras = ["datafusion"], marker = "sys_platform != 'emscripten' and sys_platform != 'win32'", specifier = ">=0.1.0rc2,<0.2" },
+    { name = "einfold", extras = ["datafusion"], marker = "sys_platform != 'emscripten' and sys_platform != 'win32'", specifier = ">=0.1.0rc2,<0.3" },
     { name = "jax", marker = "platform_machine != 'x86_64' or sys_platform != 'darwin'" },
     { name = "maturin", specifier = ">=1.9.1" },
     { name = "py-spy", specifier = ">=0.4.0" },
diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py
index d0ff2c8f..eb12555a 100644
--- a/xarray_sql/ad.py
+++ b/xarray_sql/ad.py
@@ -43,7 +43,11 @@
 
 [sql_all][xarray_sql.ad.sql_all] runs several statements that take
 ``grad`` of the same loss for the price of one backward pass, which is
-how a training step updates each of a model's parameter tables. For the
+how a training step updates each of a model's parameter tables. For repeated
+training steps, [prepare][xarray_sql.ad.prepare] retains ddx's
+programs between calls and materializes each statement's result.
+[Context.prepare][xarray_sql.Context.prepare] handles one such statement;
+``ad.prepare`` handles several statements sharing a loss. For the
 programs underneath (``grad``, ``vjp``, ``jvp`` and ``run``, to build a
 backward pass once and run it every step), use ``ddxdb.ad`` directly;
 they take a ``Context`` like any other ``SessionContext``.
@@ -52,12 +56,13 @@
 from __future__ import annotations
 
 import re
+import threading
 from collections.abc import Sequence
 
 from datafusion import SessionContext
 from ddxdb import ad as _ddx
 
-__all__ = ["STOP_GRADIENT", "calls_ad", "register", "sql", "sql_all"]
+__all__ = ["STOP_GRADIENT", "calls_ad", "prepare", "register", "sql", "sql_all"]
 
 STOP_GRADIENT = _ddx.STOP_GRADIENT
 """The SQL name of the stop-gradient function, ``ddx_stop_gradient``."""
@@ -138,3 +143,44 @@ def sql_all(ctx: SessionContext, statements: Sequence[str]) -> list:
     frames: list = _ddx.sql_all(ctx, list(statements))
     wrap = getattr(ctx, "_wrap", None)
     return [wrap(f) for f in frames] if wrap else frames
+
+
+def prepare(ctx: SessionContext, statements: Sequence[str]):
+    """Prepare SQL gradient statements for repeated execution.
+
+    ddx retains the differentiation programs and checks table schemas,
+    views and session settings on each call. This avoids rebuilding those
+    programs in a training loop while allowing tables to be registered
+    again. Several statements can share a loss, as with ``sql_all``::
+
+        step = xql.ad.prepare(ctx, [
+            f"WITH loss AS ({LOSS}) "
+            "SELECT i, val - 0.01 * grad(loss, val) AS val FROM w ORDER BY i"
+        ])
+        [rows] = step()
+
+    On a ``Context``, a call may first write Dataset values into registered
+    tables, using the same validation as ``Context.prepare``::
+
+        [rows] = step(w=new_weights)
+
+    Returns a callable producing one materialized ``pyarrow.Table`` per
+    statement. Results remain valid after later calls. Writes and execution
+    are serialized for calls to this prepared function. Other users of the
+    context must coordinate concurrent writes themselves.
+    """
+    if isinstance(statements, str):
+        raise TypeError("statements is a list of SQL statements")
+    run = _ddx.prepare(ctx, list(statements))
+    lock = threading.Lock()
+
+    def execute(**tables):
+        with lock:
+            if tables:
+                write = getattr(ctx, "_write", None)
+                if write is None:
+                    raise TypeError("Dataset writes require xarray_sql.Context")
+                write(tables)
+            return [frame.to_arrow_table() for frame in run()]
+
+    return execute
diff --git a/xarray_sql/mlx.py b/xarray_sql/mlx.py
new file mode 100644
index 00000000..c81134b4
--- /dev/null
+++ b/xarray_sql/mlx.py
@@ -0,0 +1,258 @@
+"""Optional Apple silicon execution for SQL tensor computations.
+
+[EinFold](https://github.com/xqlsystems/einfold) recognizes tensor operations
+in DataFusion plans. This extension supplies its validated dense matrices to
+[MLX](https://github.com/ml-explore/mlx), Apple's array framework. SQL parsing,
+row existence, NULL handling and [ddx](https://github.com/xqlsystems/ddx)
+gradient construction remain in their original engines.
+
+Build the native hook with ``python accelerators/einfold_mlx/build.py``.
+Then install it on a context before planning queries::
+
+    # Launch Python with MLX_ENABLE_TF32=0 before importing MLX.
+    from xarray_sql import Context
+    from xarray_sql.mlx import MLX
+
+    accelerator = MLX(precision="float32")
+    ctx = Context().with_extensions(accelerator)
+
+``float32`` explicitly permits reduced precision for dense arithmetic on
+the GPU, including arithmetic in ddx's forward and backward SQL programs.
+The default ``float64`` uses NumPy's CPU matrix multiplication and EinFold's
+CPU attention; MLX's Metal GPU does not implement double precision. Neither
+mode replaces ddx with MLX autodiff. Unsupported shapes and exceptional
+values retain EinFold's CPU implementation. No table values are cached.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+from collections import Counter
+from typing import Literal
+
+import mlx.core as mx
+import numpy as np
+import pyarrow as pa
+from einfold.datafusion import EinFold
+
+__all__ = ["MLX"]
+
+
+class MLX:
+    """EinFold extension with an optional MLX GPU arithmetic backend.
+
+    Args:
+        precision: ``float64`` keeps double precision on the CPU.
+            ``float32`` permits rounding inputs and dense results to single
+            precision on the GPU. Results retain their SQL Arrow dtype.
+        min_ops: Minimum scalar products before crossing the native/Python
+            boundary. Smaller computations use the original CPU kernel.
+        min_gpu_ops: Minimum scalar products for GPU execution. Smaller
+            dense products use float64 CPU matrix multiplication; smaller
+            attention uses EinFold's CPU tiles.
+        max_working_bytes: Upper bound on estimated extra backend buffers.
+            Larger computations retain the CPU kernel. This is separate
+            from DataFusion's own memory reservation.
+        query_block: Maximum queries in an attention block. Bounds mask
+            storage independently of the number of queries.
+
+    The native hook is supplied by the pinned feature build in
+    ``accelerators/einfold_mlx``. An ordinary EinFold installation raises
+    an actionable error instead of silently claiming GPU execution.
+    """
+
+    def __init__(
+        self,
+        *,
+        precision: Literal["float64", "float32"] = "float64",
+        min_ops: int = 1 << 20,
+        min_gpu_ops: int = 1 << 33,
+        max_working_bytes: int = 512 << 20,
+        query_block: int = 1024,
+    ) -> None:
+        if precision not in ("float64", "float32"):
+            raise ValueError("precision must be 'float64' or 'float32'")
+        if precision == "float32" and os.environ.get("MLX_ENABLE_TF32") != "0":
+            raise ValueError(
+                "Launch Python with MLX_ENABLE_TF32=0 for full float32 "
+                "arithmetic; MLX reads this process setting on first use"
+            )
+        if (
+            min_ops < 0
+            or min_gpu_ops < 0
+            or max_working_bytes <= 0
+            or not 1 <= query_block <= 1024
+        ):
+            raise ValueError("invalid operation, memory or block limit")
+        self.precision = precision
+        self._backend = _Backend(
+            precision, min_ops, min_gpu_ops, max_working_bytes, query_block
+        )
+        try:
+            self._extension = EinFold(
+                dense_backend=self._backend.run,
+                dense_min_ops=min_ops,
+                dense_max_bytes=max_working_bytes,
+            )
+        except TypeError as exc:
+            raise ImportError(
+                "EinFold needs the native dense backend hook; run "
+                "'.venv/bin/python accelerators/einfold_mlx/build.py' "
+                "from the xarray-sql checkout"
+            ) from exc
+
+    def __datafusion_session_components__(self, ctx):
+        return self._extension.__datafusion_session_components__(ctx)
+
+    def __datafusion_session_planner__(self, ctx, fallback):
+        return self._extension.__datafusion_session_planner__(ctx, fallback)
+
+    def statistics(self) -> dict[str, int]:
+        """Completed CPU/GPU kernels and declined backend requests.
+
+        Counts describe execution, including repeated collections, rather
+        than the presence of a node in an explained plan.
+        """
+        return self._backend.statistics()
+
+
+class _Backend:
+    def __init__(
+        self,
+        precision: Literal["float64", "float32"],
+        min_ops: int,
+        min_gpu_ops: int,
+        max_working_bytes: int,
+        query_block: int,
+    ) -> None:
+        self.precision = precision
+        self.min_ops = min_ops
+        self.min_gpu_ops = min_gpu_ops
+        self.max_working_bytes = max_working_bytes
+        self.query_block = query_block
+        self._counts: Counter[str] = Counter()
+        self._lock = threading.Lock()
+
+    def statistics(self) -> dict[str, int]:
+        with self._lock:
+            return dict(self._counts)
+
+    def _count(self, key: str) -> None:
+        with self._lock:
+            self._counts[key] += 1
+
+    def _decline(self):
+        self._count("declined")
+
+    def run(self, kind, columns, shape, scale, ranges):
+        # The native side has already checked positional keys, duplicates,
+        # SQL NULLs and the domain. It also refuses non-finite factors.
+        w, x, z, y = shape
+        ops = w * x * z * y if kind == "matmul" else w * z * (x + y)
+        cells = sum(len(c) for c in columns)
+        output = w * x * y if kind == "matmul" else w * y
+        if ops < self.min_ops or (cells + output) * 32 > self.max_working_bytes:
+            return self._decline()
+        arrays = [c.to_numpy(zero_copy_only=True) for c in columns]
+        if self.precision == "float64" or ops < self.min_gpu_ops:
+            if kind != "matmul":
+                return self._decline()
+            a, b = arrays
+            result = np.matmul(a.reshape(w, x, z), b.reshape(w, z, y))
+            device = "cpu"
+        else:
+            # Keep overflow, underflow and exceptional SQL arithmetic on
+            # the CPU. In particular, masked infinities must stay absent
+            # rows instead of participating in a zero-weighted product.
+            maximum = []
+            for a in arrays:
+                magnitudes = np.abs(a)
+                hi = magnitudes.max(initial=0.0)
+                if (
+                    not np.isfinite(hi)
+                    or hi > 1e15
+                    or np.any((magnitudes > 0) & (magnitudes < 1e-15))
+                ):
+                    return self._decline()
+                maximum.append(hi)
+            if kind == "attention":
+                mode, constant = scale
+                multiplier = constant if mode == 1 else 1.0
+                if mode == 2:
+                    if constant == 0:
+                        return self._decline()
+                    multiplier = 1.0 / constant
+                # Scores too close to single precision's range, including
+                # infinity, keep the original double precision algorithm.
+                bound = maximum[0] * maximum[1] * x * abs(multiplier)
+                if not np.isfinite(bound) or bound > 1e30:
+                    return self._decline()
+            block_size = self.query_block
+            if kind == "attention":
+                available = self.max_working_bytes - (cells + output) * 32
+                block_size = min(block_size, available // (z * 16))
+                if block_size < 1:
+                    return self._decline()
+            with mx.stream(mx.gpu):
+                inputs = [mx.array(a.astype(np.float32)) for a in arrays]
+                if kind == "matmul":
+                    a, b = inputs
+                    result = mx.matmul(a.reshape(w, x, z), b.reshape(w, z, y))
+                    mx.eval(result)
+                    result = np.asarray(result).astype(np.float64)
+                else:
+                    q, kt, v = inputs
+                    q, k, v = (
+                        q.reshape(w, x),
+                        kt.reshape(x, z).T,
+                        v.reshape(z, y),
+                    )
+                    # Query tiles bound the mask and temporary score storage.
+                    # Ranges encode all of EinFold's normalized key masks,
+                    # including strict, offset and sliding-window masks.
+                    result = np.zeros((w, y), dtype=np.float64)
+                    for start in range(0, w, block_size):
+                        stop = min(w, start + block_size)
+                        lo, hi = np.asarray(ranges[start:stop]).T
+                        nonempty = lo < hi
+                        if not nonempty.any():
+                            continue
+                        key_start = int(lo[nonempty].min())
+                        key_stop = int(hi[nonempty].max())
+                        key_count = key_stop - key_start
+                        causal_hi = np.clip(
+                            np.arange(stop - start)
+                            + key_count
+                            - (stop - start)
+                            + 1,
+                            0,
+                            key_count,
+                        )
+                        if np.all(lo == key_start) and np.all(hi == key_stop):
+                            allowed = None
+                        elif np.all(lo == key_start) and np.array_equal(
+                            hi - key_start, causal_hi
+                        ):
+                            # MLX aligns a causal mask at the lower right.
+                            # Its last query then sees this slice's last key.
+                            allowed = "causal"
+                        else:
+                            keys = mx.arange(key_start, key_stop)[None, :]
+                            allowed = (keys >= mx.array(lo[:, None])) & (
+                                keys < mx.array(hi[:, None])
+                            )
+                        block = mx.fast.scaled_dot_product_attention(
+                            q[None, None, start:stop, :],
+                            k[None, None, key_start:key_stop, :],
+                            v[None, None, key_start:key_stop, :],
+                            scale=multiplier,
+                            mask=allowed,
+                        )
+                        mx.eval(block)
+                        result[start:stop] = np.asarray(block)[0, 0]
+            device = "gpu"
+        if not np.isfinite(result).all():
+            return self._decline()
+        self._count(f"{device}_{kind}")
+        return pa.array(result.reshape(-1), type=pa.float64())
diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py
index 8ec88bcf..db7c82c9 100644
--- a/xarray_sql/reader.py
+++ b/xarray_sql/reader.py
@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import itertools
+import math
 import os
 import sys
 from collections.abc import Callable, Iterable, Iterator
@@ -21,9 +22,9 @@
 import xarray as xr
 
 from .df import (
+    DEFAULT_BATCH_SIZE,
     Block,
     Chunks,
-    DEFAULT_BATCH_SIZE,
     _block_len,
     _block_metadata,
     _block_slices_from_resolved,
@@ -296,11 +297,10 @@ def _in_memory_partitions(
     # isn't contiguous in memory, and flattening it would copy. More
     # partitions than cores only adds overhead, since an in-memory table has
     # no partitions to prune.
+    # Count the Cartesian chunk grid without constructing slices that this
+    # in-memory path never reads, even for very finely chunked inputs.
     count = min(
-        sum(
-            1
-            for _ in _block_slices_from_resolved(ds, resolve_chunks(ds, chunks))
-        ),
+        math.prod(len(axis) for axis in resolve_chunks(ds, chunks).values()),
         os.cpu_count() or 1,
     )
     # `per_partition` batches of `size` rows each make `count` partitions,
@@ -324,7 +324,7 @@ def read_xarray_table(
     _iteration_callback: (
         Callable[[Block, list[str] | None], None] | None
     ) = None,
-) -> "LazyArrowStreamTable":
+) -> LazyArrowStreamTable:
     """Create a lazy DataFusion table from an xarray Dataset.
 
     This is the simplest way to register xarray data with DataFusion.
diff --git a/xarray_sql/shared.py b/xarray_sql/shared.py
new file mode 100644
index 00000000..40376cf7
--- /dev/null
+++ b/xarray_sql/shared.py
@@ -0,0 +1,708 @@
+"""Explicit shared-stage execution for deterministic SQL training graphs.
+
+[EinFold](https://github.com/xqlsystems/einfold) recognizes array operations
+inside DataFusion plans. This executor materializes reused common table
+expressions (CTEs) between those plans. An installed
+[MLX](https://github.com/ml-explore/mlx) extension can execute the numerical
+kernels on Apple's GPU. Stage boundaries currently exchange Arrow tables.
+
+[ddx](https://github.com/xqlsystems/ddx) differentiates the original SQL,
+including virtual view definitions, before choosing its own materialization
+boundaries. Forward staging never replaces the query ddx differentiates.
+
+This experimental API accepts SELECTs with flat, nonrecursive WITH clauses
+and the deterministic numerical functions listed below. It requires the
+native feature build in ``accelerators/einfold_mlx``. Tables may be updated through
+``Context._write`` or registered again between calls; calls on this executor
+are serialized. Coordinate/schema changes rebuild physical plans.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+import uuid
+from collections import Counter, OrderedDict
+from collections.abc import Sequence
+
+import einfold._native as _einfold
+import numpy as np
+import pyarrow as pa
+from datafusion import SessionContext
+from ddxdb import UnsupportedExpression
+from ddxdb import ad as _ddx
+from ddxdb import prepare as _prepare_ad
+from ddxdb._ddxdb import to_sql as _to_sql
+
+from xarray_sql._native import _SortedArrowTable
+from xarray_sql._native import shared_ctes as _shared_ctes
+from xarray_sql.ad import calls_ad
+from xarray_sql.sql import _copy, _tables_read
+
+__all__ = ["SharedExecutor"]
+
+_FUNCTIONS = frozenset(
+    [
+        "abs",
+        "avg",
+        "ceil",
+        "coalesce",
+        "count",
+        "ddx_stop_gradient",
+        "exp",
+        "erf",
+        "floor",
+        "greatest",
+        "least",
+        "ln",
+        "log",
+        "max",
+        "min",
+        "nullif",
+        "pow",
+        "power",
+        "round",
+        "sign",
+        "sin",
+        "cos",
+        "sqrt",
+        "sum",
+        "tanh",
+        "row_number",
+    ]
+)
+_DYNAMIC_FILTERS = tuple(
+    f"datafusion.optimizer.enable{kind}_dynamic_filter_pushdown"
+    for kind in ("", "_join", "_topk", "_aggregate")
+)
+
+
+def _quote(name):
+    return '"' + name.replace('"', '""') + '"'
+
+
+def _snapshot(rows):
+    return pa.Table.from_arrays(
+        [
+            pa.concat_arrays(c.chunks)
+            if c.chunks
+            else pa.array([], type=c.type)
+            for c in rows.columns
+        ],
+        schema=rows.schema,
+    )
+
+
+def _canonical(rows, keys):
+    """Avoid sorting the already ordered positional output of dense kernels."""
+    if rows.num_rows < 2:
+        return rows
+    tied = None
+    for name, _ in keys:
+        column = rows[name]
+        if column.null_count or not pa.types.is_integer(column.type):
+            return rows.sort_by(keys)
+        values = column.to_numpy(zero_copy_only=False)
+        descending = values[1:] < values[:-1]
+        if tied is not None:
+            descending &= tied
+        if descending.any():
+            return rows.sort_by(keys)
+        equal = values[1:] == values[:-1]
+        tied = equal if tied is None else tied & equal
+        if not np.any(tied):
+            return rows
+    return rows
+
+
+def _join_only(plan):
+    """Join rows can remain logical until a consumer reduces them."""
+    joined = False
+    stack = [plan]
+    allowed = {
+        "Join",
+        "CrossJoin",
+        "Projection",
+        "Filter",
+        "SubqueryAlias",
+        "TableScan",
+    }
+    while stack:
+        node = stack.pop()
+        kind = type(node.to_variant()).__name__
+        if kind not in allowed:
+            return False
+        joined |= kind in {"Join", "CrossJoin"}
+        stack.extend(node.inputs())
+    return joined
+
+
+class _Compiled:
+    def __init__(self, owner, frame, *, deferred=False, graph=None):
+        self.frame, self.schema = frame, frame.schema()
+        self.deferred = deferred
+        self.graph = graph
+        self.plan = None
+        self.epoch = owner._epoch
+        self.batch = None
+        self.inputs = {
+            name: owner._versions.get(name, 0)
+            for name in _tables_read(frame.logical_plan())
+        }
+        if not deferred and graph is None:
+            self._physical(owner)
+
+    def _physical(self, owner):
+        start = time.perf_counter()
+        self.plan = self.frame.execution_plan()
+        owner._counts["physical_plans"] += 1
+        owner._times["planning_seconds"] += time.perf_counter() - start
+
+    def valid(self, owner):
+        return self.epoch == owner._epoch and all(
+            owner._versions.get(name, 0) == generation
+            for name, generation in self.inputs.items()
+        )
+
+    def run(self, owner):
+        if self.graph is not None:
+            # ddx recreates transient view providers each pass. Rebind the
+            # compound stage to those providers while retaining its buffers.
+            return self.graph.run(rebind=True)
+        if self.plan is None:
+            self._physical(owner)
+        start = time.perf_counter()
+        plan = _copy(owner.ctx, self.plan)
+        batches = [
+            b.to_pyarrow()
+            for p in range(plan.partition_count)
+            for b in owner.ctx.execute(plan, p)
+        ]
+        # Physical simplification can prove a boolean nonnullable while
+        # the logical schema still permits NULL. Keep the declared schema
+        # when collecting, without copying or converting column values.
+        batches = [
+            pa.RecordBatch.from_arrays(b.columns, schema=self.schema)
+            for b in batches
+        ]
+        owner._counts["plan_executions"] += 1
+        owner._times["execution_seconds"] += time.perf_counter() - start
+        return pa.Table.from_batches(batches, self.schema)
+
+    def store(self, owner, name, rows):
+        started = time.perf_counter()
+        keys = [
+            (f.name, "ascending")
+            for f in rows.schema
+            if pa.types.is_integer(f.type) or pa.types.is_string(f.type)
+        ]
+        if keys:
+            rows = _canonical(rows, keys)
+        arrays = [c.combine_chunks() for c in rows.columns]
+        compatible = (
+            self.batch is not None
+            and self.batch.schema == rows.schema
+            and self.batch.num_rows == rows.num_rows
+        )
+        if compatible and self.batch is not None:
+            # Integer/string keys must retain both their domain and order:
+            # cached scan statistics were derived from these coordinates.
+            compatible = all(
+                new.null_count == 0 and old.null_count == 0
+                if pa.types.is_floating(new.type)
+                else new.equals(old)
+                for new, old in zip(arrays, self.batch.columns)
+            )
+        if compatible and self.batch is not None:
+            for new, old in zip(arrays, self.batch.columns):
+                if pa.types.is_floating(new.type):
+                    width = new.type.byte_width
+                    a = memoryview(new.buffers()[1])[
+                        new.offset * width : (new.offset + len(new)) * width
+                    ]
+                    memoryview(old.buffers()[1])[:] = a
+            owner._counts["stage_buffer_writes"] += 1
+        else:
+            if self.batch is not None:
+                if self.batch.num_rows != rows.num_rows:
+                    reason = "rows"
+                elif self.batch.schema != rows.schema:
+                    reason = "schema"
+                else:
+                    reason = "keys_or_nulls"
+                owner._counts[f"stage_replacements_{reason}"] += 1
+                owner._versions[name] = owner._versions.get(name, 0) + 1
+            # concat_arrays always owns its value buffers, including a
+            # single input chunk. No stage overwrites its producer's data.
+            self.batch = pa.RecordBatch.from_arrays(
+                [owner._intern_key(a) for a in arrays], schema=rows.schema
+            )
+            owner._counts["stage_buffer_allocations"] += 1
+        owner.ctx.deregister_table(name)
+        if keys:
+            owner.ctx.register_table(
+                name, _SortedArrowTable(self.batch, [k for k, _ in keys])
+            )
+        else:
+            owner.ctx.register_record_batches(name, [[self.batch]])
+        owner._times["stage_storage_seconds"] += time.perf_counter() - started
+
+
+class _Backend(_ddx.DataFusionBackend):
+    def __init__(self, owner):
+        super().__init__(owner.ctx)
+        self.owner = owner
+        self._catalog: list = []
+        self._changes = 0
+        self._deferred: dict[str, _Compiled] = {}
+
+    def select_all(self, name):
+        deferred = self._deferred.get(name)
+        if deferred is None:
+            return super().select_all(name)
+        # ddx asks for a named Read with the saved table's schema. A view
+        # would expand to its entire definition in Substrait. Serialize a
+        # schema-only table, then restore the view before any execution.
+        empty = pa.RecordBatch.from_arrays(
+            [pa.array([], type=f.type) for f in deferred.schema],
+            schema=deferred.schema,
+        )
+        self.ctx.deregister_table(name)
+        try:
+            self.ctx.register_record_batches(name, [[empty]])
+            return super().select_all(name)
+        finally:
+            self.ctx.deregister_table(name)
+            self.ctx.register_view(name, deferred.frame)
+
+    def drop_table(self, name):
+        self._deferred.pop(name, None)
+        super().drop_table(name)
+
+    def materialize(self, name, plan):
+        super().materialize(name, plan)
+        self.owner._versions[name] = self.owner._versions.get(name, 0) + 1
+
+    def view_sql(self, name):
+        if name in self.owner._views:
+            return self.owner._views[name]
+        return super().view_sql(name)
+
+    def _bind_views(self, sql):
+        # ddx synthesizes nested WITH queries while extracting shared CTEs.
+        # Rejecting those probes makes it abandon staging and inline the
+        # entire model. The restriction on user forward graphs is separate.
+        parts = _shared_ctes(sql, allow_nested=True)
+        defined = {name for name, *_ in parts}
+        needed = {
+            ref
+            for _, _, refs, _ in parts
+            for ref in refs
+            if ref in self.owner._views and ref not in defined
+        }
+        if not needed:
+            return sql
+        # ddx extracts reused CTEs from large views. A simple view may be
+        # left to the engine; bind its original SELECT in that case too.
+        definitions = ", ".join(
+            f"{_quote(name)} AS ({self.owner._views[name]})"
+            for name in sorted(needed)
+        )
+        return f'WITH {definitions} SELECT * FROM ({sql}) AS "__xql_original"'
+
+    def plan_sql(self, sql):
+        return super().plan_sql(self._bind_views(sql))
+
+    def plan_key(self, sql):
+        return super().plan_key(self._bind_views(sql))
+
+    def session_key(self):
+        # Internal shared stages may be created while preparing a program.
+        # They are execution state, rather than changes to its public inputs.
+        try:
+            settings = (
+                SessionContext.sql(
+                    self.ctx,
+                    "SELECT name,value FROM information_schema.df_settings ORDER BY name",
+                )
+                .to_arrow_table()
+                .to_pylist()
+            )
+        except ValueError:
+            return None
+        catalog = []
+        for c in sorted(self.ctx.catalog_names()):
+            for s in sorted(self.ctx.catalog(c).schema_names()):
+                if s == "information_schema":
+                    continue
+                schema = self.ctx.catalog(c).schema(s)
+                for t in sorted(schema.table_names()):
+                    if t.startswith(("__ddx_", "__xql_shared_")):
+                        continue
+                    table = schema.table(t)
+                    path = ".".join(_quote(p) for p in (c, s, t))
+                    value = (
+                        self.ctx.table(path).logical_plan()
+                        if table.kind == "view"
+                        else table.schema
+                    )
+                    catalog.append((path, value))
+        if catalog != self._catalog:
+            self._catalog, self._changes = catalog, self._changes + 1
+        return repr((self._changes, settings, self.owner._views))
+
+    def compile(self, plan):
+        # DataFusion 55's Substrait consumer applies Read projections only
+        # to TableScan roots; a deferred view expands to a join instead.
+        # ddx's SQL writer preserves the same plan and handles those reads.
+        frame = (
+            SessionContext.sql(self.ctx, _to_sql(plan, "datafusion"))
+            if self._deferred
+            else _ddx._dataframe(self.ctx, plan)
+        )
+        stack = [frame.logical_plan()]
+        unions = 0
+        while stack:
+            node = stack.pop()
+            inputs = node.inputs()
+            if type(node.to_variant()).__name__ == "Union":
+                unions += max(0, len(inputs) - 1)
+            stack.extend(inputs)
+        if unions >= 32:
+            # Large cotangent sums have many independently reduced UNION
+            # branches. Materialize those reductions sequentially; running
+            # all their join build sides together can exhaust memory.
+            graph = self._stage_graph(plan, reductions=True)
+            if graph is not None:
+                self.owner._counts["union_reduction_graphs"] += 1
+                return _Compiled(self.owner, frame, graph=graph)
+        deferred = _join_only(frame.logical_plan())
+        if deferred:
+            try:
+                parts = _shared_ctes(
+                    _to_sql(plan, "datafusion"), allow_nested=True
+                )
+                deferred = all(set(fns) <= _FUNCTIONS for _, _, _, fns in parts)
+            except (UnsupportedExpression, ValueError):
+                deferred = False
+        return _Compiled(self.owner, frame, deferred=deferred)
+
+    def _stage_graph(self, plan, **options):
+        try:
+            sql = _to_sql(plan, "datafusion")
+            return _Graph(self.owner, sql, **options)
+        except (UnsupportedExpression, ValueError):
+            return None
+
+    def _refresh_compiled(self, plan, compiled):
+        if compiled.valid(self.owner):
+            return False
+        fresh = self.compile(plan)
+        batch = compiled.batch
+        if (
+            compiled.graph is not None
+            and fresh.graph is not None
+            and compiled.graph.sql == fresh.graph.sql
+        ):
+            fresh.graph = compiled.graph
+        compiled.__dict__.update(fresh.__dict__)
+        compiled.batch = batch
+        return True
+
+    def materialize_compiled(self, name, plan, compiled):
+        changed = self._refresh_compiled(plan, compiled)
+        if compiled.deferred:
+            if changed:
+                self.owner._versions[name] = (
+                    self.owner._versions.get(name, 0) + 1
+                )
+            self._deferred[name] = compiled
+            self.ctx.deregister_table(name)
+            self.ctx.register_view(name, compiled.frame)
+            self.owner._counts["deferred_stages"] += 1
+            return
+        rows = compiled.run(self.owner)
+        compiled.store(self.owner, name, rows)
+        compiled.epoch = self.owner._epoch
+
+    def returns_rows_compiled(self, plan, compiled):
+        self._refresh_compiled(plan, compiled)
+        return compiled.run(self.owner).num_rows > 0
+
+
+class _Graph:
+    def __init__(self, owner, sql, *, reductions=False):
+        self.owner = owner
+        self.reductions = reductions
+        self.sql = sql
+        self.parts = _shared_ctes(sql)
+        self.defs = {name: (query, refs) for name, query, refs, _ in self.parts}
+        if len(self.defs) != len(self.parts):
+            raise ValueError("duplicate CTE names cannot be staged")
+        for _, _, _, functions in self.parts:
+            unknown = set(functions) - _FUNCTIONS
+            if unknown:
+                raise ValueError(
+                    "shared execution requires deterministic numerical SQL; "
+                    f"unsupported functions: {', '.join(sorted(unknown))}"
+                )
+        needed = set()
+
+        def visit(name, stack):
+            if name in stack:
+                raise ValueError("cyclic CTE dependencies cannot be staged")
+            if name in needed:
+                return
+            needed.add(name)
+            for dep in self.defs[name][1]:
+                if dep in self.defs:
+                    visit(dep, stack | {name})
+
+        visit("", set())
+        counts = Counter(
+            dep
+            for name in needed
+            for dep in self.defs[name][1]
+            if dep in needed
+        )
+        self.boundaries = {
+            name
+            for name, count in counts.items()
+            if count > 1 and name not in owner.inline_ctes
+        }
+        if reductions:
+            self.boundaries.update(
+                name
+                for name, _, _, functions in self.parts
+                if name
+                and name in needed
+                and set(functions) & {"sum", "avg", "min", "max", "count"}
+            )
+        self.names = {
+            name: owner._prefix + uuid.uuid4().hex
+            for name in self.boundaries | {""}
+        }
+        self.order = [name for name, *_ in self.parts if name in self.names]
+        self.compiled = {}
+        self.tables = {
+            dep
+            for name in needed
+            for dep in self.defs[name][1]
+            if dep not in self.defs
+        }
+
+    def _sql(self, target, aliases):
+        definitions = dict(aliases)
+
+        def include(name):
+            if name in definitions:
+                return
+            if name in self.boundaries and name != target:
+                definitions[name] = f"SELECT * FROM {_quote(self.names[name])}"
+                return
+            for dep in self.defs[name][1]:
+                if dep in self.defs:
+                    include(dep)
+            definitions[name] = self.defs[name][0]
+
+        for dep in self.defs[target][1]:
+            if dep in self.defs:
+                include(dep)
+        prefix = ", ".join(
+            f"{_quote(name)} AS ({query})"
+            for name, query in definitions.items()
+        )
+        return (f"WITH {prefix} " if prefix else "") + self.defs[target][0]
+
+    def run(self, stack=(), *, rebind=False):
+        owner = self.owner
+        aliases = {}
+        stages = 0
+        for name in sorted(self.tables & owner._views.keys()):
+            if name in stack:
+                raise ValueError("cyclic virtual views cannot be staged")
+            graph = owner._view_graph(name)
+            rows = graph.run((*stack, name))
+            stages += graph.stages
+            graph.compiled[""].store(owner, graph.names[""], rows)
+            aliases[name] = f"SELECT * FROM {_quote(graph.names[''])}"
+        for name in self.order:
+            compiled = self.compiled.get(name)
+            if rebind or compiled is None or not compiled.valid(owner):
+                frame = SessionContext.sql(owner.ctx, self._sql(name, aliases))
+                fresh = _Compiled(owner, frame)
+                if compiled is not None:
+                    fresh.batch = compiled.batch
+                compiled = self.compiled[name] = fresh
+            rows = compiled.run(owner)
+            stages += 1
+            if name:
+                compiled.store(owner, self.names[name], rows)
+            compiled.epoch = owner._epoch
+        self.stages = stages
+        if self.reductions:
+            owner._counts["reduction_stage_executions"] += stages
+        else:
+            owner._counts["forward_stages"] = stages
+        return rows
+
+
+class SharedExecutor:
+    """Cache shared SQL stages and ddx programs over live parameter tables.
+
+    Install EinFold or MLX on ``ctx`` before constructing this executor.
+    Register model SQL with ``register_view``; it stays as SQL until ddx
+    has differentiated it. ``prepare`` returns a function producing an
+    Arrow table, and ``prepare_all`` produces a list of Arrow tables.
+
+    ``inline_ctes`` names CTEs to keep inside consumers, for example an
+    attention score/softmax chain that EinFold can fuse as one operator.
+    Only reused, reachable CTEs outside that set become stage boundaries.
+
+    Physical plans retain input buffers, with fresh join state each run.
+    Re-registering a public table invalidates those plans. Returned results
+    are independent snapshots. Changes to virtual views require preparing
+    forward queries again. Do not modify context settings while executing.
+    """
+
+    def __init__(
+        self,
+        ctx,
+        *,
+        inline_ctes: Sequence[str] = (),
+        max_key_cache_bytes: int = 256 << 20,
+    ):
+        if not getattr(_einfold, "RESET_ON_DECODE", False):
+            raise ImportError(
+                "rebuild accelerators/einfold_mlx for safe plan reuse"
+            )
+        self.ctx = ctx
+        if max_key_cache_bytes < 0:
+            raise ValueError("max_key_cache_bytes must be nonnegative")
+        self._max_key_bytes = max_key_cache_bytes
+        self._key_bytes = 0
+        self._keys: OrderedDict[tuple, list] = OrderedDict()
+        self.inline_ctes = frozenset(inline_ctes)
+        self._views: dict[str, str] = {}
+        self._graphs: dict[str, _Graph] = {}
+        self._counts: Counter[str] = Counter()
+        self._times: Counter[str] = Counter()
+        self._epoch = 0
+        self._versions: dict[str, int] = {}
+        self._prefix = "__xql_shared_" + uuid.uuid4().hex + "_"
+        self._lock = threading.RLock()
+        self._catalog: tuple | None = None
+        for key in _DYNAMIC_FILTERS:
+            SessionContext.sql(ctx, f"SET {key} = false").collect()
+        self._backend = _Backend(self)
+
+    def _intern_key(self, array):
+        size = array.nbytes
+        if (
+            not (
+                pa.types.is_integer(array.type)
+                or pa.types.is_string(array.type)
+            )
+            or size > self._max_key_bytes
+        ):
+            return pa.concat_arrays([array])
+        # The fingerprint is only an index: equality checks every value and
+        # NULL before sharing a buffer. Owned integer/string buffers are
+        # immutable; floating point values always get separate storage.
+        fingerprint = (
+            str(array.type),
+            len(array),
+            array.null_count,
+            tuple(
+                array[i].as_py() for i in (0, len(array) // 2, len(array) - 1)
+            )
+            if len(array)
+            else (),
+        )
+        candidates = self._keys.get(fingerprint, [])
+        for old in candidates:
+            if array.equals(old):
+                self._keys.move_to_end(fingerprint)
+                self._counts["key_cache_hits"] += 1
+                self._counts["key_copy_bytes_avoided"] += size
+                return old
+        owned = pa.concat_arrays([array])
+        while self._keys and self._key_bytes + size > self._max_key_bytes:
+            _, evicted = self._keys.popitem(last=False)
+            self._key_bytes -= sum(a.nbytes for a in evicted)
+        # Eviction may have removed this fingerprint's bucket.
+        self._keys.setdefault(fingerprint, []).append(owned)
+        self._keys.move_to_end(fingerprint)
+        self._key_bytes += size
+        return owned
+
+    def register_view(self, name: str, sql: str):
+        """Keep a deterministic SELECT as a virtual view for execution and ddx."""
+        with self._lock:
+            _Graph(self, sql)  # Validate before changing the view catalog.
+            if self._views.get(name) != sql:
+                self._views[name] = sql
+                self._graphs.clear()
+                self._epoch += 1
+
+    def _view_graph(self, name):
+        if name not in self._graphs:
+            self._graphs[name] = _Graph(self, self._views[name])
+        return self._graphs[name]
+
+    def _refresh(self):
+        catalog = tuple(
+            sorted(
+                (name, generation)
+                for name, generation in self.ctx._deregistrations.items()
+                if not name.startswith(("__ddx_", "__xql_shared_"))
+            )
+        )
+        if catalog != self._catalog:
+            self._catalog = catalog
+            self._epoch += 1
+
+    def prepare(self, sql: str):
+        """Prepare one SELECT; differentiation always uses its original SQL."""
+        if calls_ad(sql):
+            run = self.prepare_all([sql])
+            return lambda **tables: run(**tables)[0]
+        graph = _Graph(self, sql)
+
+        def execute(**tables):
+            with self._lock:
+                self.ctx._write(tables)
+                self._refresh()
+                return _snapshot(graph.run())
+
+        return execute
+
+    def prepare_all(self, statements: Sequence[str]):
+        """Prepare ddx statements sharing a backward pass over original views."""
+        prepared = _prepare_ad(self._backend, statements)
+        finals: dict[str, _Compiled] = {}
+
+        def execute(**tables):
+            with self._lock:
+                self.ctx._write(tables)
+                self._refresh()
+                with prepared.run() as rewritten:
+                    results = []
+                    for sql in rewritten.statements:
+                        compiled = finals.get(sql)
+                        if compiled is None or not compiled.valid(self):
+                            compiled = finals[sql] = _Compiled(
+                                self, SessionContext.sql(self.ctx, sql)
+                            )
+                        results.append(_snapshot(compiled.run(self)))
+                    return results
+
+        return execute
+
+    def statistics(self):
+        """Executed stages, physical-plan builds and cumulative timings."""
+        return (
+            dict(self._counts)
+            | dict(self._times)
+            | {"key_cache_bytes": self._key_bytes}
+        )
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index f64166a4..bf801903 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -1,6 +1,8 @@
 import re
 import time
 import uuid
+from collections.abc import Callable
+from types import ModuleType
 
 import numpy as np
 import pyarrow as pa
@@ -9,7 +11,6 @@
 from datafusion.catalog import Schema
 from datafusion.expr import TableScan
 from datafusion.plan import ExecutionPlan, LogicalPlan
-from types import ModuleType
 
 from . import cftime as cft
 from .df import (
@@ -38,14 +39,30 @@
 except ImportError:  # pragma: no cover - depends on the environment
     _ad = None
 
+_DYNAMIC_FILTERS = tuple(
+    f"datafusion.optimizer.enable{kind}_dynamic_filter_pushdown"
+    for kind in ("", "_join", "_topk", "_aggregate")
+)
+
 
 class Context(SessionContext):
     """A datafusion `SessionContext` that also supports `xarray.Dataset`s.
 
     Also available under its original name, ``XarrayContext``.
+    ``threads=`` sets DataFusion's target parallel partitions, overriding
+    that setting in an explicit ``config``. The default is the CPU count.
+    It is a concurrency target, rather than a fixed number of OS threads.
     """
 
-    def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
+    def __init__(
+        self,
+        config: SessionConfig | None = None,
+        *args,
+        threads: int | None = None,
+        **kwargs,
+    ):
+        if threads is not None and threads < 1:
+            raise ValueError("threads must be positive")
         if config is None:
             # Split a query's work across partitions only where enough rows
             # reach it, by DataFusion's estimate. DataFusion otherwise splits
@@ -65,6 +82,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
                     "true",
                 )
             )
+        if threads is not None:
+            config = config.with_target_partitions(threads)
         super().__init__(config, *args, **kwargs)
         # Track registered xarray Datasets so XarrayDataFrame can recover
         # defaults (dimension_columns) and metadata (var/dataset attrs,
@@ -78,6 +97,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
         # registered again from Python was deregistered first; SQL that
         # replaces or drops a table counts too (see sql()).
         self._deregistrations: dict[str, int] = {}
+        self._custom_udfs: set[str] = set()
+        self._trusted_math_udfs: set[str] = set()
         # With pyproj installed, every context speaks CRS out of the box:
         # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform.
         if _proj is not None:
@@ -87,6 +108,41 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
         if _ad is not None:
             _ad.register(self)
 
+    def register_udf(self, udf):
+        """Register a scalar function and track overrides for optional lowering."""
+        result = super().register_udf(udf)
+        self._custom_udfs.add(udf.name.lower())
+        self._trusted_math_udfs.discard(udf.name.lower())
+        return result
+
+    def register_udaf(self, udaf):
+        """Track aggregate-function overrides for optional numerical lowering."""
+        result = super().register_udaf(udaf)
+        self._custom_udfs.add(udaf.name.lower())
+        self._trusted_math_udfs.discard(udaf.name.lower())
+        return result
+
+    def register_udwf(self, udwf):
+        """Track window-function overrides for optional numerical lowering."""
+        result = super().register_udwf(udwf)
+        self._custom_udfs.add(udwf.name.lower())
+        self._trusted_math_udfs.discard(udwf.name.lower())
+        return result
+
+    def with_extensions(self, *extensions) -> "Context":
+        """Install DataFusion extension bundles, retaining xarray support.
+
+        Returns another handle on the same session, with this Context's
+        Dataset metadata and prepared-query invalidation state shared.
+        Extension codecs belong to the returned handle, so use it for
+        subsequent queries, as with DataFusion's own method.
+        """
+        session = super().with_extensions(*extensions)
+        derived = object.__new__(type(self))
+        derived.__dict__.update(self.__dict__)
+        derived.ctx = session.ctx
+        return derived
+
     def from_dataset(
         self,
         name: str,
@@ -214,9 +270,13 @@ def _from_dataset(
             table = read_xarray_table(
                 input_table, chunks, coord_arrays=coord_arrays
             )
+            register(table_name, table)
+        elif schema is None:
+            # Uniform Datasets can register their shared Arrow buffers
+            # directly, without a temporary catalog entry and lookup.
+            self.register_record_batches(table_name, partitions)
         else:
-            table = self._in_memory_table(partitions)
-        register(table_name, table)
+            register(table_name, self._in_memory_table(partitions))
         self._maybe_register_cftime_udf(input_table)
         return self
 
@@ -324,16 +384,11 @@ def prepare(self, query: str) -> "PreparedQuery":
             A [PreparedQuery][xarray_sql.PreparedQuery]; call it to run
             the query.
 
-        Raises:
-            NotImplementedError: If the query calls ``grad``, ``vjp`` or
-                ``jvp``, which only ``sql()`` computes for now.
+        Queries calling ``grad``, ``vjp`` or ``jvp`` retain ddx's programs
+        through [xarray_sql.ad.prepare][]. Their first call prepares the
+        programs; later calls check the catalog and settings and refresh
+        them if necessary.
         """
-        if _ad is not None and _ad.calls_ad(query):
-            # TODO(#294): prepare the query's programs with ddx.
-            raise NotImplementedError(
-                "prepare() doesn't compute grad, vjp or jvp yet; run the "
-                "query with sql(), or several with xarray_sql.ad.sql_all()"
-            )
         return PreparedQuery(self, query)
 
     def _write(self, writes: dict[str, xr.Dataset]) -> None:
@@ -415,10 +470,19 @@ class PreparedQuery:
     query again. A call raises if a table the query reads has been
     registered again since it was prepared. Functions such as ``now()``
     are evaluated once, when the query is prepared.
+
+    A query calling ``grad``, ``vjp`` or ``jvp`` instead retains ddx's
+    programs. Its first call prepares them, and subsequent calls refresh
+    them when schemas, views or settings change. Tables may be registered
+    again under the same names, as with [xarray_sql.ad.prepare][].
     """
 
     def __init__(self, ctx: Context, query: str):
         self._ctx = ctx
+        self._gradient: Callable[..., list[pa.Table]] | None = None
+        if _ad is not None and _ad.calls_ad(query):
+            self._gradient = _ad.prepare(ctx, [query])
+            return
         start = time.perf_counter()
         df = SessionContext.sql(ctx, query)
         plan = df.execution_plan()
@@ -441,11 +505,13 @@ def __init__(self, ctx: Context, query: str):
         }
         # A plan keeps state from a run (a hash join keeps the table it
         # built), so each call runs a fresh copy of it, decoded from its
-        # serialized form. Copying a plan copies its in-memory tables, so
-        # over large ones it's slower than planning again; a plan that reads
-        # a lazy table, or is nested deeper than protobuf decodes, can't be
-        # copied at all. Either way, each call plans the query again.
+        # serialized form. The default codec copies in-memory tables, so
+        # over large ones it can be slower than planning again; extensions
+        # can supply a codec that resets state while retaining input buffers.
+        # A plan the session's codecs cannot copy is planned again each call.
         self._plan: ExecutionPlan | None = None
+        if not _safe_physical_reuse(ctx):
+            return
         start = time.perf_counter()
         try:
             _copy(ctx, plan)
@@ -464,6 +530,9 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table:
         Returns:
             The query's rows.
         """
+        if self._gradient is not None:
+            (rows,) = self._gradient(**tables)
+            return rows
         for name in tables:
             if name not in self._registered:
                 raise ValueError(
@@ -495,7 +564,22 @@ def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan:
     # TODO(#283): reset the plan's state with datafusion-python's binding of
     # reset_plan_states, once it has one, rather than decode it again.
     # `plan` itself never runs, so it has no state to leave behind.
-    return ExecutionPlan.from_bytes(ctx, plan.to_bytes())
+    return ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx))
+
+
+def _safe_physical_reuse(ctx: SessionContext) -> bool:
+    """Dynamic filters retain value-dependent bounds even after a reset."""
+    names = ", ".join(repr(name) for name in _DYNAMIC_FILTERS)
+    try:
+        rows = SessionContext.sql(
+            ctx,
+            f"SELECT value FROM information_schema.df_settings WHERE name IN ({names})",
+        ).to_arrow_table()
+    except ValueError:
+        return False
+    return rows.num_rows == len(_DYNAMIC_FILTERS) and all(
+        value == "false" for value in rows["value"].to_pylist()
+    )
 
 
 def _tables_read(plan: LogicalPlan) -> set[str]: