diff --git a/AGENTS.md b/AGENTS.md
index fc419a15..dfb638b5 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -8,7 +8,16 @@ summarizes recurring maintainer review feedback so changes land clean.
 - Keep docstrings and comments self-contained. Do **not** put GitHub issue or PR
   numbers in docstrings or code comments; a reader should not need the issue
   tracker to understand the code. Issue references belong in the commit message
-  and PR description (e.g. `Closes #189`), not in the source.
+  and PR description (e.g. `Closes #189`), not in the source. The one exception
+  is a workaround's `TODO`, which names the issue that would remove it:
+  `TODO(#NN): <action item>`, or `TODO(owner/repo#NN): <action item>` for an
+  issue in another repository. The comment still explains the workaround.
+- Introduce every reference before relying on it. The first time a document
+  (a docstring, comment, README or docs page) mentions a project, paper,
+  algorithm, dataset or term of art, such as einfold, FlashAttention or a
+  section of another project's design, say in a few words what it is, and link
+  it where there is a link. Don't assume the reader has the context of the
+  conversation that produced the code.
 - Do not reference the review conversation, chat, or "the reporter" in comments.
   Describe the behavior, not how it came up.
 
diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py
new file mode 100644
index 00000000..ecd59dfb
--- /dev/null
+++ b/benchmarks/ml/01_attention.py
@@ -0,0 +1,270 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "xarray-sql[ml,duckdb]",
+#   "jax",
+# ]
+#
+# [tool.uv.sources]
+# xarray-sql = { path = "../../", editable = true }
+# ///
+"""Demo 1: causal attention in SQL, on several engines, with its gradient.
+
+Attention, the core of a transformer, is ``softmax(Q·Kᵀ/√d)·V``: each
+token's output is an average of the value vectors ``V``, weighted by how
+well the token's query in ``Q`` matches each earlier token's key in ``K``.
+Here ``Q``, ``K`` and ``V`` are tables of ``(t, d, x)`` (token, feature,
+value), registered from xarray Datasets, and attention is two sums over
+joins: the scores, a sum of ``Q`` times ``K`` over features, and the
+average of ``V`` weighted by ``exp(score)``.
+
+This SQL is the naive plan: it stores every one of the N×N scores.
+FlashAttention (Dao et al., 2022, https://arxiv.org/abs/2205.14135)
+computes the same result without storing them, by streaming blocks of
+keys past blocks of queries and keeping a running softmax. einfold
+(https://github.com/xqlsystems/einfold), a query optimizer that rewrites
+SQL plans into array algorithms like this one, aims to find that plan
+from this SQL, in both of the spellings below, with no rule written for
+attention. It is the first of einfold's three target demos
+(``docs/demos.md`` in its repository). This script is the baseline that
+plan must match and then beat:
+
+* **forward**: each spelling of the model on each *host* (einfold's word for an engine
+  that runs a plan; here DataFusion and DuckDB), checked against the same
+  attention in JAX, Google's library for differentiable array programs;
+* **gradient**: a loss that reads the model, and its gradient with
+  respect to every value of Q, K and V, as one SQL statement. ddx (https://github.com/xqlsystems/ddx),
+  which differentiates SQL queries, computes it, and ``jax.grad`` checks
+  it;
+* **masks are absent rows**: an infinity in V reaches only the queries
+  allowed to see it, as SQL's semantics require of any faster plan;
+* **a control that is not attention**: the Nadaraya–Watson estimator, a
+  Gaussian-kernel weighted average used in regridding. Any rule that
+  speeds up attention must speed up this query too.
+
+Run ``uv run benchmarks/ml/01_attention.py --help`` for the options.
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import xarray as xr
+from _bench import Report, assert_close, cli, hosts
+
+import xarray_sql as xql
+
+jax.config.update("jax_enable_x64", True)
+
+D = 64  # features per head; the scores are scaled by 1/√64 = 1/8.
+
+# The model and its loss are separate, as in a neural network written in
+# Python (compare attention() and loss() below). The model is a view, a
+# named query: each spelling of attention defines the view ``attention``,
+# and the forward pass and the loss read it as if it were a table.
+
+# Attention in einfold's window spelling. The causal mask is a join
+# predicate: a masked (i, j) pair is a row that does not exist. Subtracting
+# each row's max only keeps exp() from overflowing: it cancels out of the
+# ratio w / z, so its gradient sums to zero, and ddx differentiates it
+# like any other window.
+ATTENTION_WINDOW = """
+CREATE OR REPLACE VIEW attention AS
+WITH scores AS (
+  SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s
+  FROM q JOIN k ON q.d = k.d
+  WHERE q.t >= k.t
+  GROUP BY q.t, k.t),
+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 AS t, v.d, SUM(w.w / n.z * v.x) AS x
+FROM weights w
+JOIN v ON w.j = v.t
+JOIN norm n ON w.i = n.i
+GROUP BY w.i, v.d
+"""
+
+# The same attention as ddx spells it: each row's max as a GROUP BY,
+# joined back to the scores.
+ATTENTION_GROUP_BY = """
+CREATE OR REPLACE VIEW attention AS
+WITH scores AS (
+  SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s
+  FROM q JOIN k ON q.d = k.d
+  WHERE q.t >= k.t
+  GROUP BY q.t, k.t),
+row_max AS (
+  SELECT i, MAX(s) AS m FROM scores GROUP BY i),
+weights AS (
+  SELECT s.i, s.j, exp(s.s - r.m) AS w
+  FROM scores s JOIN row_max r ON s.i = r.i),
+norm AS (
+  SELECT i, SUM(w) AS z FROM weights GROUP BY i)
+SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x
+FROM weights w
+JOIN v ON w.j = v.t
+JOIN norm n ON w.i = n.i
+GROUP BY w.i, v.d
+"""
+
+SPELLINGS = {"window": ATTENTION_WINDOW, "group by": ATTENTION_GROUP_BY}
+
+# The forward pass: run the model.
+FORWARD = "SELECT t, d, x FROM attention ORDER BY t, d"
+
+# The loss, half the output's squared norm, calls the model by reading the
+# view. The SELECT asks for its gradient with respect to every value of Q,
+# K and V. The three grad() calls differentiate one loss, so ddx runs one
+# backward pass for all three, as jax.grad(loss, argnums=(0, 1, 2)) does.
+GRADIENT = """
+WITH loss AS (
+  SELECT SUM(x * x) / 2 AS l FROM attention)
+SELECT q.t, q.d,
+       grad(loss, q.x) AS dq,
+       grad(loss, k.x) AS dk,
+       grad(loss, v.x) AS dv
+FROM q
+JOIN k ON q.t = k.t AND q.d = k.d
+JOIN v ON q.t = v.t AND q.d = v.d
+ORDER BY q.t, q.d
+"""
+
+# Nadaraya–Watson: at each target point, the average of the sources' y
+# weighted by a Gaussian kernel of their distance (bandwidth h = 0.5).
+# Structurally attention: a weighted average with weights exp(score).
+NADARAYA_WATSON = """
+WITH kernel AS (
+  SELECT t.p, s.q, exp(-(power(t.x - s.x, 2) + power(t.y - s.y, 2)) / 0.25) AS w
+  FROM target t CROSS JOIN source s)
+SELECT k.p, SUM(k.w * s.v) / SUM(k.w) AS v
+FROM kernel k JOIN source s ON k.q = s.q
+GROUP BY k.p
+ORDER BY k.p
+"""
+
+
+def attention(q, k, v):
+    """The reference: causal softmax attention, as JAX writes it."""
+    n = q.shape[0]
+    scores = q @ k.T / np.sqrt(D)
+    causal = jnp.tril(jnp.ones((n, n), dtype=bool))
+    weights = jax.nn.softmax(jnp.where(causal, scores, -jnp.inf), axis=-1)
+    return weights @ v
+
+
+def loss(q, k, v):
+    return jnp.sum(attention(q, k, v) ** 2) / 2
+
+
+def table(values: np.ndarray) -> xr.Dataset:
+    """A (t, d) matrix as the XQL table ``(t, d, x)``."""
+    t, d = values.shape
+    return xr.Dataset(
+        {"x": (("t", "d"), values)},
+        coords={"t": np.arange(t), "d": np.arange(d)},
+    )
+
+
+def matrix(rows, name: str = "x") -> np.ndarray:
+    """The rows of a ``(t, d, …)`` result as a (t, d) matrix."""
+    return xql.to_dataset(rows, dims=["t", "d"])[name].values
+
+
+def main():
+    parser = cli(__doc__.splitlines()[0], sizes=[64, 256], size_name="tokens")
+    args = parser.parse_args()
+    report = Report("attention", "tokens", args.reps)
+    rng = np.random.default_rng(0)
+
+    for n in args.sizes:
+        q, k, v = (rng.normal(size=(n, D)) for _ in range(3))
+        want = attention(q, k, v)
+        want_grads = jax.grad(loss, argnums=(0, 1, 2))(q, k, v)
+
+        for host in hosts(args.hosts):
+            for name, data in zip("qkv", (q, k, v)):
+                host.load(name, table(data))
+
+            for spelling, model in SPELLINGS.items():
+                host.execute(model)
+                rows = report.time(
+                    f"forward ({spelling})", host, n, host.query, FORWARD
+                )
+                report.check(assert_close("forward", matrix(rows), want))
+                report.note("operators", host.operators(FORWARD))
+
+                step = f"gradient ({spelling})"
+                if not host.differentiates:
+                    report.skip(step, host, n, "no ddx backend")
+                    continue
+                rows = report.time(step, host, n, host.query, GRADIENT)
+                errors = [
+                    assert_close(f"d/d{name}", matrix(rows, f"d{name}"), ref)
+                    for name, ref in zip("qkv", want_grads)
+                ]
+                report.check(max(errors))
+
+    check_masks_are_absent_rows(args.hosts)
+    check_nadaraya_watson(args.hosts, report)
+    report.finish(args.csv)
+
+
+def check_masks_are_absent_rows(host_names):
+    """An infinity in V at token j reaches the queries at or after j, and
+    no others: a masked pair is skipped, not multiplied by a zero weight,
+    which would turn 0·∞ into NaN. (xarray-sql reads NaN as missing data,
+    a NULL that SUM skips, so an infinity is the value that tells.)"""
+    n, j = 16, 9
+    rng = np.random.default_rng(1)
+    q, k, v = (rng.normal(size=(n, D)) for _ in range(3))
+    v[j, 0] = np.inf
+    print(f"\n  masks: +inf in V at token {j}, feature 0")
+    for host in hosts(host_names):
+        for name, data in zip("qkv", (q, k, v)):
+            host.load(name, table(data))
+        for spelling, model in SPELLINGS.items():
+            host.execute(model)
+            out = matrix(host.query(FORWARD))
+            assert np.isfinite(out[:j]).all(), (host.name, spelling)
+            assert np.isposinf(out[j:, 0]).all(), (host.name, spelling)
+            assert np.isfinite(out[:, 1:]).all(), (host.name, spelling)
+            print(f"  {spelling:<9} {host.name:<11} ok")
+
+
+def check_nadaraya_watson(host_names, report):
+    """The control: a kernel-weighted average that is not attention."""
+    rng = np.random.default_rng(2)
+    targets, sources = rng.uniform(size=(64, 2)), rng.uniform(size=(256, 2))
+    y = np.sin(4 * sources[:, 0]) + np.cos(3 * sources[:, 1])
+    dist2 = ((targets[:, None, :] - sources[None, :, :]) ** 2).sum(-1)
+    w = np.exp(-dist2 / 0.25)
+    want = (w @ y) / w.sum(axis=1)
+
+    target = xr.Dataset(
+        {"x": ("p", targets[:, 0]), "y": ("p", targets[:, 1])},
+        coords={"p": np.arange(len(targets))},
+    )
+    source = xr.Dataset(
+        {"x": ("q", sources[:, 0]), "y": ("q", sources[:, 1]), "v": ("q", y)},
+        coords={"q": np.arange(len(sources))},
+    )
+    for host in hosts(host_names):
+        host.load("target", target)
+        host.load("source", source)
+        rows = report.time(
+            f"nadaraya-watson ({len(sources)} sources)",
+            host,
+            len(targets),
+            host.query,
+            NADARAYA_WATSON,
+        )
+        got = xql.to_dataset(rows, dims=["p"])["v"].values
+        report.check(assert_close("nadaraya-watson", got, want))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py
new file mode 100644
index 00000000..ae670d29
--- /dev/null
+++ b/benchmarks/ml/02_nanogpt.py
@@ -0,0 +1,623 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "xarray-sql[ml,duckdb]",
+#   "jax",
+# ]
+#
+# [tool.uv.sources]
+# xarray-sql = { path = "../../", editable = true }
+# ///
+"""Demo 2: a GPT in SQL, trained by gradient descent.
+
+nanoGPT (https://github.com/karpathy/nanoGPT) is Andrej Karpathy's small,
+readable implementation of GPT-2, the language model. This is its model,
+``model.py``, as one SQL query: the cross-entropy loss of a batch of
+characters from Shakespeare, computed by a transformer whose parameters
+are tables. ddx (https://github.com/xqlsystems/ddx), which differentiates
+SQL queries, takes the loss's gradient with respect to every parameter,
+so a training step is SQL too.
+
+=================  ============================================
+Model piece        In SQL
+=================  ============================================
+Embeddings         a join of the batch with ``wte`` and ``wpe``
+Layer norm         per-position ``AVG``s, joined back, times a gain
+Attention          per head and causal, as in ``01_attention.py``
+MLP                two einsums around GELU
+Output layer       an einsum with ``wte`` again (tied, as in GPT-2)
+Loss               log-sum-exp per position, minus the target's logit
+Gradients          ddx: ``grad(loss, w)``, one statement per table
+Optimizer step     ``w - lr * grad(loss, w)``, the new table; or AdamW,
+                   with Adam's moments as columns beside ``w`` and
+                   the learning rate schedule as a table
+=================  ============================================
+
+As in nanoGPT, the model and the loss are separate. The model is the
+view ``logits``; the loss is SQL that reads it, as nanoGPT's
+``F.cross_entropy`` reads the model's logits.
+
+Each training step updates every parameter table with its own
+statement; the statements share one backward pass. By default the
+optimizer is SGD at a fixed learning rate, the simplest step to read.
+``--optimizer adamw`` trains as nanoGPT does instead: it stores the
+gradients as tables, then updates each parameter table with AdamW. The same model in JAX
+(Google's library for differentiable array programs), from the same
+initial parameters and the same batches, is the reference: its loss at
+every step, and its parameters after every step, must match the SQL's.
+
+This is Part A of the second target demo of einfold
+(https://github.com/xqlsystems/einfold), a query optimizer that rewrites
+SQL plans into array algorithms; the demo is described in einfold's
+``docs/demos.md``. Part B builds on it: it writes the tricks behind a
+record of the NanoGPT speedrun (a public race to train GPT-2 fastest)
+as a small diff to this SQL.
+
+Differences from nanoGPT's ``model.py`` and ``train.py``:
+
+* **Size.** nanoGPT's Shakespeare model has 6 layers, 6 heads, 384
+  channels and 256 characters of context, 10.65M parameters. This one has
+  1 layer, 2 heads, 16 channels and 16 characters, 4,416 parameters (the
+  constants below).
+* **GELU** is its tanh approximation, as GPT-2 computed it. nanoGPT's
+  ``nn.GELU()`` uses the exact form, through the error function ``erf``,
+  which DataFusion doesn't have. The two differ by less than 5e-4.
+* **The optimizer**, by default, is SGD at a fixed learning rate.
+  ``--optimizer adamw`` matches nanoGPT's: AdamW with its Shakespeare
+  settings, weight decay on all but the layer norms, gradients clipped to
+  a norm of 1, and the learning rate warming up, then decaying along a
+  cosine. The schedule is stretched to ``--steps``, where nanoGPT's spans
+  5,000. The JAX reference for this step matches ``torch.optim.AdamW``
+  after ``clip_grad_norm_`` exactly.
+* **No dropout**, nanoGPT's default, so that a run is deterministic and
+  can be checked. Its Shakespeare configuration trains with 0.2.
+* **Q, K and V are three tables.** nanoGPT's ``c_attn`` holds the same
+  parameters as one matrix, split after the product.
+
+Like nanoGPT's default (``bias = False`` in ``train.py``), the model has
+no biases, in linear layers or layer norms.
+
+Run ``uv run benchmarks/ml/02_nanogpt.py --help`` for the options.
+"""
+
+from __future__ import annotations
+
+import pathlib
+import urllib.request
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import xarray as xr
+from _bench import Report, assert_close, cli, hosts
+
+import xarray_sql as xql
+
+jax.config.update("jax_enable_x64", True)
+
+# A tiny character-level GPT. The comments name each constant in nanoGPT's
+# GPTConfig (https://github.com/karpathy/nanoGPT/blob/master/model.py) and
+# its value in nanoGPT's Shakespeare model
+# (https://github.com/karpathy/nanoGPT/blob/master/config/train_shakespeare_char.py).
+LAYERS = 1  # n_layer = 6: transformer blocks
+EMBED = 16  # n_embd = 384; e: channels of the residual stream
+HEADS = 2  # n_head = 6; h: attention heads
+HEAD = EMBED // HEADS  # c: channels per head
+HIDDEN = 4 * EMBED  # f: channels of the MLP, 4 × n_embd as in nanoGPT
+CONTEXT = 16  # block_size = 256; pos: characters per sequence
+LEARNING_RATE = 0.5  # for SGD; nanoGPT's AdamW uses 1e-3, below
+
+# AdamW (``--optimizer adamw``), with the settings of nanoGPT's Shakespeare
+# model: train.py's defaults, and config/train_shakespeare_char.py's.
+MAX_LR = 1e-3  # learning_rate
+MIN_LR = 1e-4  # min_lr
+WARMUP = 100 / 5000  # warmup_iters, as a fraction of max_iters
+BETA1 = 0.9  # beta1
+BETA2 = 0.99  # beta2
+EPS = 1e-8  # torch.optim.AdamW's default
+WEIGHT_DECAY = 0.1  # weight_decay, on every table but the layer norms' gains
+GRAD_CLIP = 1.0  # grad_clip: the largest norm of all gradients together
+
+# Each parameter table's dimensions; its values are the column w. In
+# nanoGPT, wq, wk and wv are attn.c_attn; wo is attn.c_proj; w_fc and
+# w_proj are mlp.c_fc and mlp.c_proj; ln_1, ln_2 and ln_f are the layer
+# norms' weights.
+PARAMS = {
+    "wte": ("tok", "e"),
+    "wpe": ("pos", "e"),
+    "ln_1": ("layer", "e"),
+    "wq": ("layer", "e", "h", "c"),
+    "wk": ("layer", "e", "h", "c"),
+    "wv": ("layer", "e", "h", "c"),
+    "wo": ("layer", "h", "c", "e"),
+    "ln_2": ("layer", "e"),
+    "w_fc": ("layer", "e", "f"),
+    "w_proj": ("layer", "f", "e"),
+    "ln_f": ("e",),
+}
+
+# --------------------------------------------------------------------------
+# The model, in SQL. Activations are tables of (seq, pos, e, val): one row
+# per sequence in the batch, position, and channel.
+# --------------------------------------------------------------------------
+
+
+def layer_norm(x: str, out: str, gain: str, layer: int | None = None) -> str:
+    """Normalize each position of ``x`` to zero mean and unit variance,
+    then scale each channel by the table ``gain`` (of ``layer``, if set)."""
+    of_layer = "" if layer is None else f" AND g.layer = {layer}"
+    return f"""
+{out}_mean AS (
+  SELECT seq, pos, AVG(val) AS mean FROM {x} GROUP BY seq, pos),
+{out}_var AS (
+  SELECT x.seq, x.pos, AVG(power(x.val - m.mean, 2)) AS var
+  FROM {x} x JOIN {out}_mean m ON x.seq = m.seq AND x.pos = m.pos
+  GROUP BY x.seq, x.pos),
+{out} AS (
+  SELECT x.seq, x.pos, x.e, (x.val - m.mean) / sqrt(v.var + 1e-5) * g.w AS val
+  FROM {x} x
+  JOIN {out}_mean m ON x.seq = m.seq AND x.pos = m.pos
+  JOIN {out}_var v ON x.seq = v.seq AND x.pos = v.pos
+  JOIN {gain} g ON x.e = g.e{of_layer}),"""
+
+
+# GELU's tanh approximation: x/2 · (1 + tanh(√(2/π) · (x + 0.044715 x³))).
+GELU_SCALE = np.sqrt(2 / np.pi)
+
+
+def block(layer: int, x: str, out: str) -> str:
+    """One transformer block: ``out = x + attention(ln(x))``, then
+    ``out = out + mlp(ln(out))``."""
+    l = f"l{layer}"
+    projections = "".join(
+        f"""
+{l}_{name} AS (
+  SELECT x.seq, x.pos, p.h, p.c, SUM(x.val * p.w) AS val
+  FROM {l}_ln1 x JOIN w{name} p ON x.e = p.e
+  WHERE p.layer = {layer}
+  GROUP BY x.seq, x.pos, p.h, p.c),"""
+        for name in ("q", "k", "v")
+    )
+    return f"""{layer_norm(x, f"{l}_ln1", "ln_1", layer)}{projections}
+{l}_scores AS (
+  SELECT q.seq, q.h, q.pos AS i, k.pos AS j,
+         SUM(q.val * k.val) / {np.sqrt(HEAD)} AS s
+  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),
+{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
+  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
+  GROUP BY a.seq, a.i, a.h, v.c),
+{l}_attn 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),
+{l}_mid AS (
+  SELECT x.seq, x.pos, x.e, x.val + a.val AS val
+  FROM {x} x JOIN {l}_attn a ON x.seq = a.seq AND x.pos = a.pos AND x.e = a.e),
+{layer_norm(f"{l}_mid", f"{l}_ln2", "ln_2", layer)}
+{l}_fc AS (
+  SELECT x.seq, x.pos, p.f, SUM(x.val * p.w) AS val
+  FROM {l}_ln2 x JOIN w_fc p ON x.e = p.e
+  WHERE p.layer = {layer}
+  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
+  FROM {l}_fc),
+{l}_mlp 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),
+{out} AS (
+  SELECT x.seq, x.pos, x.e, x.val + m.val AS val
+  FROM {l}_mid x JOIN {l}_mlp m ON x.seq = m.seq AND x.pos = m.pos AND x.e = m.e),"""
+
+
+def model() -> 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)
+    )
+    return f"""
+CREATE OR REPLACE VIEW logits AS
+WITH x0 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),
+{blocks}
+{layer_norm(f"x{LAYERS}", "final", "ln_f").rstrip(",")}
+SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val
+FROM final x JOIN wte te ON x.e = te.e
+GROUP BY x.seq, x.pos, te.tok
+"""
+
+
+# The loss, separate from the model, as in nanoGPT: the mean cross-entropy
+# of the model's logits against each next character of the batch.
+# Subtracting each row's max only keeps exp() from overflowing: it cancels
+# out of the log-sum-exp.
+#
+# The loss reads the model's view once, into the CTE ``predicted``, and
+# uses that twice: a host expands a view at every read, but DuckDB
+# computes a CTE that is read twice only once.
+LOSS = """
+predicted AS (
+  SELECT seq, pos, tok, val FROM logits),
+shifted AS (
+  SELECT seq, pos, tok, val,
+         MAX(val) OVER (PARTITION BY seq, pos) AS m
+  FROM predicted),
+log_sum_exp AS (
+  SELECT seq, pos, MAX(m) + ln(SUM(exp(val - m))) AS lse
+  FROM shifted GROUP BY seq, pos),
+loss AS (
+  SELECT AVG(s.lse - l.val) AS l
+  FROM log_sum_exp s
+  JOIN batch b ON s.seq = b.seq AND s.pos = b.pos
+  JOIN predicted l ON l.seq = b.seq AND l.pos = b.pos AND l.tok = b.next)"""
+
+
+def step_statements(lr: float) -> list[str]:
+    """One SGD step: each parameter table's next version, as a statement."""
+    return [
+        f"WITH {LOSS} SELECT {', '.join(dims)}, "
+        f"w - {lr} * grad(loss, w) AS w FROM {table} "
+        f"ORDER BY {', '.join(dims)}"
+        for table, dims in PARAMS.items()
+    ]
+
+
+def gradient_statements() -> list[str]:
+    """The gradient of the loss with respect to each parameter table, as
+    the table ``{name}_grad``'s rows. The statements share one backward
+    pass."""
+    return [
+        f"WITH {LOSS} SELECT {', '.join(dims)}, grad(loss, w) AS g "
+        f"FROM {table} ORDER BY {', '.join(dims)}"
+        for table, dims in PARAMS.items()
+    ]
+
+
+def adamw_statement(table: str, step: int) -> str:
+    """One AdamW step on a parameter table with columns ``w``, ``m`` and
+    ``v`` (the weight and Adam's two moments), from its gradient in
+    ``{table}_grad`` and the learning rate in ``schedule``. This is
+    ``torch.optim.AdamW``, after nanoGPT clips the gradients."""
+    dims = PARAMS[table]
+    keys = ", ".join(f"p.{d}" for d in dims)
+    same_row = " AND ".join(f"p.{d} = g.{d}" for d in dims)
+    decay = 0.0 if table.startswith("ln_") else WEIGHT_DECAY
+    every_grad = "\n    UNION ALL ".join(
+        f"SELECT g FROM {t}_grad" for t in PARAMS
+    )
+    return f"""
+WITH grad_norm AS (
+  SELECT sqrt(SUM(g * g)) AS norm FROM (
+    {every_grad})),
+clip AS (
+  SELECT LEAST(1.0, {GRAD_CLIP} / (norm + 1e-6)) AS c FROM grad_norm),
+moments AS (
+  SELECT {keys}, p.w,
+         {BETA1} * p.m + (1 - {BETA1}) * g.g * clip.c AS m,
+         {BETA2} * p.v + (1 - {BETA2}) * power(g.g * clip.c, 2) AS v
+  FROM {table} p
+  JOIN {table}_grad g ON {same_row}
+  CROSS JOIN clip)
+SELECT {keys},
+       p.w * (1 - s.lr * {decay})
+         - s.lr * (p.m / (1 - power({BETA1}, s.step + 1)))
+           / (sqrt(p.v / (1 - power({BETA2}, s.step + 1))) + {EPS}) AS w,
+       p.m,
+       p.v
+FROM moments p JOIN schedule s ON s.step = {step}
+ORDER BY {keys}
+"""
+
+
+def loss_statement() -> str:
+    return f"WITH {LOSS} SELECT l FROM loss"
+
+
+# --------------------------------------------------------------------------
+# The same model in JAX: the reference.
+# --------------------------------------------------------------------------
+
+
+def jax_layer_norm(x, gain):
+    mean = x.mean(-1, keepdims=True)
+    var = ((x - mean) ** 2).mean(-1, keepdims=True)
+    return (x - mean) / jnp.sqrt(var + 1e-5) * gain
+
+
+def jax_loss(p, tok, nxt):
+    x = p["wte"][tok] + p["wpe"][None]
+    causal = jnp.tril(jnp.ones((CONTEXT, CONTEXT), dtype=bool))
+    for layer in range(LAYERS):
+        h = jax_layer_norm(x, p["ln_1"][layer])
+        q, k, v = (
+            jnp.einsum("bte,ehc->bhtc", h, p[w][layer])
+            for w in ("wq", "wk", "wv")
+        )
+        s = jnp.einsum("bhic,bhjc->bhij", q, k) / np.sqrt(HEAD)
+        a = jax.nn.softmax(jnp.where(causal, s, -jnp.inf), axis=-1)
+        heads = jnp.einsum("bhij,bhjc->bhic", a, v)
+        x = x + jnp.einsum("bhtc,hce->bte", heads, p["wo"][layer])
+        h = jax_layer_norm(x, p["ln_2"][layer])
+        hidden = jax.nn.gelu(h @ p["w_fc"][layer], approximate=True)
+        x = x + hidden @ p["w_proj"][layer]
+    logits = jax_layer_norm(x, p["ln_f"]) @ p["wte"].T
+    target = jnp.take_along_axis(logits, nxt[..., None], -1)[..., 0]
+    return jnp.mean(jax.nn.logsumexp(logits, -1) - target)
+
+
+def jax_adamw(params, moments, grads, step, lr):
+    """The reference AdamW step: new parameters and moments."""
+    norm = np.sqrt(sum(float(jnp.sum(g**2)) for g in grads.values()))
+    clip = min(1.0, GRAD_CLIP / (norm + 1e-6))
+    new_params, new_moments = {}, {}
+    for name, w in params.items():
+        g = grads[name] * clip
+        m, v = moments[name]
+        m = BETA1 * m + (1 - BETA1) * g
+        v = BETA2 * v + (1 - BETA2) * g**2
+        decay = 0.0 if name.startswith("ln_") else WEIGHT_DECAY
+        m_hat = m / (1 - BETA1 ** (step + 1))
+        v_hat = v / (1 - BETA2 ** (step + 1))
+        new_params[name] = w * (1 - lr * decay) - lr * m_hat / (
+            jnp.sqrt(v_hat) + EPS
+        )
+        new_moments[name] = (m, v)
+    return new_params, new_moments
+
+
+# --------------------------------------------------------------------------
+# Data and parameters, as xarray Datasets.
+# --------------------------------------------------------------------------
+
+SHAKESPEARE = (
+    "https://raw.githubusercontent.com/karpathy/char-rnn/master/"
+    "data/tinyshakespeare/input.txt"
+)
+
+
+def shakespeare() -> str:
+    """Tiny Shakespeare, cached in ~/.cache; a sonnet offline."""
+    cache = (
+        pathlib.Path.home() / ".cache" / "xarray-sql" / "tinyshakespeare.txt"
+    )
+    if not cache.exists():
+        cache.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            with urllib.request.urlopen(SHAKESPEARE, timeout=10) as r:
+                cache.write_bytes(r.read())
+        except OSError:
+            return SONNET_18 * 4
+    return cache.read_text()
+
+
+def batches(text: str, size: int, rng):
+    """Endless batches of (seq, pos) characters and the characters after
+    them, as one Dataset: the table ``batch(seq, pos, tok, next)``."""
+    vocab = sorted(set(text))
+    ids = np.array([vocab.index(c) for c in text])
+    while True:
+        starts = rng.integers(0, len(ids) - CONTEXT - 1, size)
+        tok = np.stack([ids[s : s + CONTEXT] for s in starts])
+        nxt = np.stack([ids[s + 1 : s + CONTEXT + 1] for s in starts])
+        yield xr.Dataset(
+            {"tok": (("seq", "pos"), tok), "next": (("seq", "pos"), nxt)},
+            coords={"seq": np.arange(size), "pos": np.arange(CONTEXT)},
+        )
+
+
+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,
+    }
+    params = {}
+    for name, dims in PARAMS.items():
+        shape = [sizes[d] for d in dims]
+        if name.startswith("ln_"):
+            params[name] = np.ones(shape)
+        elif name in ("wo", "w_proj"):
+            params[name] = rng.normal(0, 0.02 / np.sqrt(2 * LAYERS), shape)
+        else:
+            params[name] = rng.normal(0, 0.02, shape)
+    return params
+
+
+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."""
+    dims = PARAMS[name]
+    columns = {"w": w, "m": m, "v": v}
+    return xr.Dataset(
+        {k: (dims, x) for k, x in columns.items() if x is not None},
+        coords={d: np.arange(n) for d, n in zip(dims, w.shape)},
+    )
+
+
+def as_array(name: str, rows, column: str = "w") -> np.ndarray:
+    return xql.to_dataset(rows, dims=list(PARAMS[name]))[column].values
+
+
+def schedule(steps: int) -> xr.Dataset:
+    """nanoGPT's learning rate at each step, as the table
+    ``schedule(step, lr)``: a linear warmup, then a cosine decay to the
+    minimum, with train.py's ``max_iters`` and ``lr_decay_iters`` both
+    ``steps``."""
+    warmup = round(WARMUP * steps)
+    lr = []
+    for it in range(steps):
+        if it < warmup:
+            lr.append(MAX_LR * (it + 1) / (warmup + 1))
+        else:
+            decay = (it - warmup) / (steps - warmup)
+            coeff = 0.5 * (1 + np.cos(np.pi * decay))
+            lr.append(MIN_LR + coeff * (MAX_LR - MIN_LR))
+    return xr.Dataset({"lr": ("step", lr)}, coords={"step": np.arange(steps)})
+
+
+def adamw_step(host, step: int) -> list:
+    """One AdamW step on every parameter table: the gradients, in one
+    backward pass, registered as tables; then each table's update."""
+    gradients = host.query_all(gradient_statements())
+    for name, rows in zip(PARAMS, gradients):
+        host.load(f"{name}_grad", xql.to_dataset(rows, dims=list(PARAMS[name])))
+    return [host.query(adamw_statement(name, step)) for name in PARAMS]
+
+
+def main():
+    parser = cli(__doc__.splitlines()[0], sizes=[4], size_name="batch")
+    parser.add_argument("--steps", type=int, default=5, help="training steps")
+    parser.add_argument(
+        "--optimizer",
+        choices=["sgd", "adamw"],
+        default="sgd",
+        help="SGD at a fixed learning rate (default), or nanoGPT's AdamW",
+    )
+    args = parser.parse_args()
+    adamw = args.optimizer == "adamw"
+    lr = schedule(args.steps)
+    report = Report("nanogpt", "batch", args.reps)
+    text = shakespeare()
+    vocab = len(set(text))
+    lines = len(model().splitlines()) + len(loss_statement().splitlines())
+    count = sum(
+        v.size for v in init_params(vocab, np.random.default_rng(0)).values()
+    )
+    print(
+        f"  {LAYERS}-layer GPT, vocab {vocab}, {count:,} parameters: "
+        f"the loss is {lines} lines of SQL"
+    )
+
+    for size in args.sizes:
+        init = init_params(vocab, np.random.default_rng(0))
+        data = list(
+            zip(
+                range(args.steps), batches(text, size, np.random.default_rng(1))
+            )
+        )
+
+        for host in hosts(args.hosts):
+            reference = dict(init)
+            moments = {
+                n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items()
+            }
+            for name, values in init.items():
+                state = moments[name] if adamw else ()
+                host.load(name, as_table(name, values, *state))
+            host.load("schedule", lr)
+
+            for step, batch in data:
+                host.load("batch", batch)
+                host.execute(model())
+                tok, nxt = batch["tok"].values, batch["next"].values
+
+                rows = report.time(
+                    f"loss, step {step}",
+                    host,
+                    size,
+                    host.query,
+                    loss_statement(),
+                )
+                want = jax_loss(reference, tok, nxt)
+                report.check(
+                    assert_close("loss", rows["l"][0].as_py(), want, rtol=1e-9)
+                )
+                report.note("loss", round(float(want), 6))
+                report.note("operators", host.operators(loss_statement()))
+
+                train = f"train ({args.optimizer}), step {step}"
+                if not host.differentiates:
+                    report.skip(train, host, size, "no ddx backend")
+                    continue
+                grads = jax.grad(jax_loss)(reference, tok, nxt)
+                if adamw:
+                    updated = report.time(
+                        train, host, size, adamw_step, host, step
+                    )
+                    reference, moments = jax_adamw(
+                        reference, moments, grads, step, lr["lr"].values[step]
+                    )
+                else:
+                    updated = report.time(
+                        train,
+                        host,
+                        size,
+                        host.query_all,
+                        step_statements(LEARNING_RATE),
+                    )
+                    reference = {
+                        name: w - LEARNING_RATE * grads[name]
+                        for name, w in reference.items()
+                    }
+
+                errors = []
+                for name, rows in zip(PARAMS, updated):
+                    expected = {"w": reference[name]}
+                    if adamw:
+                        expected["m"], expected["v"] = moments[name]
+                    for column, values in expected.items():
+                        errors.append(
+                            assert_close(
+                                f"{name}.{column}, step {step}",
+                                as_array(name, rows, column),
+                                values,
+                            )
+                        )
+                    host.load(
+                        name, xql.to_dataset(rows, dims=list(PARAMS[name]))
+                    )
+                report.check(max(errors))
+
+    report.finish(args.csv)
+
+
+SONNET_18 = """Shall I compare thee to a summer's day?
+Thou art more lovely and more temperate:
+Rough winds do shake the darling buds of May,
+And summer's lease hath all too short a date;
+Sometime too hot the eye of heaven shines,
+And often is his gold complexion dimm'd;
+And every fair from fair sometime declines,
+By chance or nature's changing course untrimm'd;
+But thy eternal summer shall not fade,
+Nor lose possession of that fair thou ow'st;
+Nor shall death brag thou wander'st in his shade,
+When in eternal lines to time thou grow'st:
+So long as men can breathe or eyes can see,
+So long lives this, and this gives life to thee.
+"""
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py
new file mode 100644
index 00000000..c2704518
--- /dev/null
+++ b/benchmarks/ml/03_graphcast.py
@@ -0,0 +1,724 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "xarray-sql[ml,duckdb]",
+#   "jax",
+#   "gcsfs",
+#   "zarr>=3",
+# ]
+#
+# [tool.uv.sources]
+# xarray-sql = { path = "../../", editable = true }
+# ///
+"""Demo 3: a GraphCast-shaped weather model in SQL, over ERA5 in Zarr.
+
+GraphCast (Lam et al., 2023) forecasts the atmosphere 6 hours ahead from
+the two previous states. It is a graph neural network in three stages:
+
+* an **encoder** maps the latitude–longitude grid onto a "multi-mesh",
+  a refined icosahedron that keeps the edges of every refinement level;
+* a **processor** passes messages along the mesh's edges, round after
+  round;
+* a **decoder** maps the mesh back onto the grid, as the 6-hour change.
+
+Here it is in miniature, as SQL over ERA5, ECMWF's reanalysis of the
+atmosphere since 1940. xarray-sql reads it from the Zarr store of
+WeatherBench 2, a benchmark for weather models (or a synthetic stand-in,
+offline). Every
+graph is a table of edges ``(edge, sender, receiver)``, every state a
+table of ``(id, f, val)``, and every stage the same relational steps:
+
+=================  ============================================
+GNN step           In SQL
+=================  ============================================
+Gather             join the nodes' states to the edges, by sender
+                   and by receiver
+Concatenate        ``UNION ALL``, offsetting each part's features
+MLP                einsums with ``params`` around swish, ``x / (1 + exp(-x))``
+Aggregate          ``SUM`` of the edges' messages, ``GROUP BY`` receiver
+=================  ============================================
+
+The geometry (the mesh, and which grid cells and mesh nodes are joined)
+is computed once with numpy; the model, its inputs and its forecast are
+SQL. The script checks the forecast, and the gradient of the forecast
+error, against the same model in JAX, Google's library for differentiable
+array programs. ddx (https://github.com/xqlsystems/ddx), which
+differentiates SQL queries, computes the gradient.
+
+Then it measures what the third target demo of einfold
+(https://github.com/xqlsystems/einfold, a query optimizer that rewrites
+SQL plans into array algorithms; the demo is in its ``docs/demos.md``) is
+about: how much less a regional forecast needs than a global one. It
+reports the receptive field: how many mesh nodes and grid cells each
+stage needs for the region.
+
+Simplifications from GraphCast: random weights, a 64×32 grid and six
+variables, a mesh refined 2 times rather than 6, 2 processor rounds
+rather than 16, no biases or layer norms, edge features of 3-D
+displacement and length, decoder edges from each cell's 3 nearest mesh
+nodes, and edge states that are embedded once rather than updated.
+
+Run ``uv run benchmarks/ml/03_graphcast.py --help`` for the options.
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import xarray as xr
+from ddxdb import DdxError
+from _bench import Report, assert_close, cli, hosts
+
+jax.config.update("jax_enable_x64", True)
+
+ERA5 = (
+    "gs://weatherbench2/datasets/era5/"
+    "1959-2023_01_10-6h-64x32_equiangular_conservative.zarr"
+)
+ERA5_TABLES = {
+    ("time", "longitude", "latitude"): "surface",
+    ("time", "level", "longitude", "latitude"): "atmosphere",
+}
+# The forecast's variables, as (table, variable, pressure level).
+VARIABLES = [
+    ("surface", "2m_temperature", None),
+    ("surface", "mean_sea_level_pressure", None),
+    ("surface", "10m_u_component_of_wind", None),
+    ("surface", "10m_v_component_of_wind", None),
+    ("atmosphere", "geopotential", 500),
+    ("atmosphere", "temperature", 850),
+]
+# The two input states and the state to forecast, 6 hours apart.
+TIMES = pd.to_datetime(["2020-01-01T00", "2020-01-01T06", "2020-01-01T12"])
+
+HIDDEN = 16  # channels of every node and edge state
+INPUTS = 2 * len(VARIABLES) + 3  # two states, cos(lat), sin(lon), cos(lon)
+EDGE_FEATURES = 4  # receiver minus sender in 3-D, and the edge's length
+# A region to forecast: the contiguous United States, roughly.
+BOUNDS = {"latitude": (25, 50), "longitude": (235, 295)}
+REGION = " AND ".join(
+    f"c.{coord} BETWEEN {low} AND {high}"
+    for coord, (low, high) in BOUNDS.items()
+)
+
+# --------------------------------------------------------------------------
+# The model, in SQL.
+# --------------------------------------------------------------------------
+
+
+def mlp(out: str, x: str, net: str) -> str:
+    """A two-layer MLP over the features of each ``id`` of ``x``."""
+    return f"""
+{out}_hidden AS (
+  SELECT x.id, p.g AS f, SUM(x.val * p.w) AS val
+  FROM {x} x JOIN params p ON p.net = '{net}.0' AND x.f = p.f
+  GROUP BY x.id, p.g),
+{out} AS (
+  SELECT h.id, p.g AS f, SUM(h.val / (1 + exp(-h.val)) * p.w) AS val
+  FROM {out}_hidden h JOIN params p ON p.net = '{net}.1' AND h.f = p.f
+  GROUP BY h.id, p.g),"""
+
+
+def message_passing(out, net, edges, edge_h, senders, receivers) -> str:
+    """One round of an interaction network on the graph ``edges``: each
+    edge's message is an MLP of the edge, its sender and its receiver;
+    each receiver adds an MLP of itself and its messages' sum to itself."""
+    return f"""
+{out}_edge_in AS (
+  SELECT id, f, val FROM {edge_h}
+  UNION ALL
+  SELECT e.edge AS id, s.f + {HIDDEN} AS f, s.val
+  FROM {edges} e JOIN {senders} s ON e.sender = s.id
+  UNION ALL
+  SELECT e.edge AS id, r.f + {2 * HIDDEN} AS f, r.val
+  FROM {edges} e JOIN {receivers} r ON e.receiver = r.id),
+{mlp(f"{out}_message", f"{out}_edge_in", f"{net}.edge")}
+{out}_sum AS (
+  SELECT e.receiver AS id, m.f, SUM(m.val) AS val
+  FROM {out}_message m JOIN {edges} e ON m.id = e.edge
+  GROUP BY e.receiver, m.f),
+{out}_node_in AS (
+  SELECT id, f, val FROM {receivers}
+  UNION ALL
+  SELECT id, f + {HIDDEN} AS f, val FROM {out}_sum),
+{mlp(f"{out}_update", f"{out}_node_in", f"{net}.node")}
+{out} AS (
+  SELECT r.id, r.f, r.val + u.val AS val
+  FROM {receivers} r JOIN {out}_update u ON r.id = u.id AND r.f = u.f),"""
+
+
+def era5_state(time, first_feature: int) -> list[str]:
+    """Each variable at ``time``, normalized, as (cell, feature, value)
+    rows. Only these variables, levels and times are read from ERA5."""
+    parts = []
+    for var, (table, name, level) in enumerate(VARIABLES):
+        on_level = f" AND x.level = {level}" if level else ""
+        parts.append(
+            f"""  SELECT c.cell AS id, {first_feature + var} AS f,
+         (x."{name}" - n.mean) / n.std AS val
+  FROM era5.{table} x
+  JOIN cells c ON x.latitude = c.latitude AND x.longitude = c.longitude
+  JOIN scale n ON n.var = {var}
+  WHERE x.time = TIMESTAMP '{time}'{on_level}"""
+        )
+    return parts
+
+
+def position(table: str, key: str, first_feature: int) -> list[str]:
+    """Where each row of ``table`` is on the sphere, as 3 features."""
+    features = [
+        "cos(radians(latitude))",
+        "sin(radians(longitude))",
+        "cos(radians(longitude))",
+    ]
+    return [
+        f"  SELECT {key} AS id, {first_feature + i} AS f, {feature} AS val "
+        f"FROM {table}"
+        for i, feature in enumerate(features)
+    ]
+
+
+def union(name: str, parts: list[str]) -> str:
+    return f"\n{name} AS (\n" + "\n  UNION ALL\n".join(parts) + "),"
+
+
+def edge_features(graph: str) -> str:
+    return f"(SELECT edge AS id, f, val FROM {graph}.features)"
+
+
+def model(rounds: int) -> str:
+    """The CTEs of the whole model, ending in ``forecast``: the change of
+    each normalized variable over the next 6 hours, at each cell."""
+    n = len(VARIABLES)
+    processor = "".join(
+        message_passing(
+            f"mesh{r + 1}",
+            f"processor.{r}",
+            "multimesh.edges",
+            "multimesh_h",
+            f"mesh{r}",
+            f"mesh{r}",
+        )
+        for r in range(rounds)
+    )
+    return f"""{
+        union(
+            "grid_in",
+            era5_state(TIMES[0], 0)
+            + era5_state(TIMES[1], n)
+            + position("cells", "cell", 2 * n),
+        )
+    }
+{union("mesh_in", position("mesh", "node", 0))}
+{mlp("grid_h", "grid_in", "embed.grid")}
+{mlp("mesh_h", "mesh_in", "embed.mesh")}
+{mlp("g2m_h", edge_features("g2m"), "embed.g2m")}
+{mlp("multimesh_h", edge_features("multimesh"), "embed.multimesh")}
+{mlp("m2g_h", edge_features("m2g"), "embed.m2g")}
+{message_passing("mesh0", "encoder", "g2m.edges", "g2m_h", "grid_h", "mesh_h")}
+{processor}
+{
+        message_passing(
+            "grid_out",
+            "decoder",
+            "m2g.edges",
+            "m2g_h",
+            f"mesh{rounds}",
+            "grid_h",
+        )
+    }
+forecast AS (
+  SELECT g.id, p.g AS f, SUM(g.val * p.w) AS val
+  FROM grid_out g JOIN params p ON p.net = 'output' AND g.f = p.f
+  GROUP BY g.id, p.g)"""
+
+
+def forecast_statement(rounds: int, where: str = "TRUE") -> str:
+    """The forecast change of each variable at each cell in ``where``."""
+    return f"""WITH {model(rounds)}
+SELECT c.latitude, c.longitude, f.f AS var, f.val AS change
+FROM forecast f JOIN cells c ON f.id = c.cell
+WHERE {where}
+ORDER BY c.latitude, c.longitude, f.f"""
+
+
+def gradient_statement(rounds: int) -> str:
+    """Each weight's gradient of the forecast's mean squared error
+    against what ERA5 recorded 6 hours later."""
+    n = len(VARIABLES)
+    return f"""WITH {model(rounds)},
+{union("later", era5_state(TIMES[2], 0))}
+error AS (
+  SELECT AVG(power(f.val - (y.val - x.val), 2)) AS l
+  FROM forecast f
+  JOIN later y ON f.id = y.id AND f.f = y.f
+  JOIN grid_in x ON f.id = x.id AND f.f + {n} = x.f)
+SELECT weight, grad(error, w) AS w FROM params ORDER BY weight"""
+
+
+def receptive_field(host, rounds: int) -> list[tuple[str, int]]:
+    """How many cells and mesh nodes each stage needs for the region,
+    found by walking the edges backward from it: semi-joins in SQL."""
+
+    def count(sql):
+        return host.query(f"SELECT COUNT(*) AS n FROM ({sql}) t")["n"][
+            0
+        ].as_py()
+
+    needed = f"SELECT cell AS id FROM cells c WHERE {REGION}"
+    stages = [("decoder: grid cells forecast", count(needed))]
+    needed = (
+        f"SELECT DISTINCT sender AS id FROM m2g.edges "
+        f"WHERE receiver IN ({needed})"
+    )
+    stages.append((f"processor round {rounds}: mesh nodes", count(needed)))
+    for r in reversed(range(rounds)):
+        needed = (
+            f"SELECT DISTINCT sender AS id FROM multimesh.edges "
+            f"WHERE receiver IN ({needed}) UNION SELECT id FROM ({needed}) n"
+        )
+        label = f"processor round {r}" if r else "encoder"
+        stages.append((f"{label}: mesh nodes", count(needed)))
+    needed = (
+        f"SELECT DISTINCT sender AS id FROM g2m.edges "
+        f"WHERE receiver IN ({needed})"
+    )
+    stages.append(("encoder: grid cells read", count(needed)))
+    return stages
+
+
+# --------------------------------------------------------------------------
+# The same model in JAX: the reference.
+# --------------------------------------------------------------------------
+
+
+def jax_mlp(p, net, x):
+    h = x @ p[f"{net}.0"]
+    return (h * jax.nn.sigmoid(h)) @ p[f"{net}.1"]
+
+
+def jax_message_passing(p, net, edges, edge_h, senders, receivers):
+    sender, receiver = edges
+    edge_in = jnp.concatenate(
+        [edge_h, senders[sender], receivers[receiver]], axis=1
+    )
+    message = jax_mlp(p, f"{net}.edge", edge_in)
+    total = jax.ops.segment_sum(message, receiver, len(receivers))
+    node_in = jnp.concatenate([receivers, total], axis=1)
+    return receivers + jax_mlp(p, f"{net}.node", node_in)
+
+
+def jax_forecast(p, graph, grid_in, rounds):
+    grid_h = jax_mlp(p, "embed.grid", grid_in)
+    mesh_h = jax_mlp(p, "embed.mesh", graph["mesh_in"])
+    edge_h = {
+        g: jax_mlp(p, f"embed.{g}", graph[g]["features"])
+        for g in ("g2m", "multimesh", "m2g")
+    }
+    mesh = jax_message_passing(
+        p, "encoder", graph["g2m"]["edges"], edge_h["g2m"], grid_h, mesh_h
+    )
+    for r in range(rounds):
+        mesh = jax_message_passing(
+            p,
+            f"processor.{r}",
+            graph["multimesh"]["edges"],
+            edge_h["multimesh"],
+            mesh,
+            mesh,
+        )
+    grid_out = jax_message_passing(
+        p, "decoder", graph["m2g"]["edges"], edge_h["m2g"], mesh, grid_h
+    )
+    return grid_out @ p["output"]
+
+
+def jax_error(p, graph, grid_in, target, rounds):
+    return jnp.mean((jax_forecast(p, graph, grid_in, rounds) - target) ** 2)
+
+
+# --------------------------------------------------------------------------
+# Geometry, with numpy: the multi-mesh and the edges between it and the grid.
+# --------------------------------------------------------------------------
+
+
+def icosahedron():
+    t = (1 + 5**0.5) / 2
+    vertices = [
+        (-1, t, 0),
+        (1, t, 0),
+        (-1, -t, 0),
+        (1, -t, 0),
+        (0, -1, t),
+        (0, 1, t),
+        (0, -1, -t),
+        (0, 1, -t),
+        (t, 0, -1),
+        (t, 0, 1),
+        (-t, 0, -1),
+        (-t, 0, 1),
+    ]
+    faces = [
+        (0, 11, 5),
+        (0, 5, 1),
+        (0, 1, 7),
+        (0, 7, 10),
+        (0, 10, 11),
+        (1, 5, 9),
+        (5, 11, 4),
+        (11, 10, 2),
+        (10, 7, 6),
+        (7, 1, 8),
+        (3, 9, 4),
+        (3, 4, 2),
+        (3, 2, 6),
+        (3, 6, 8),
+        (3, 8, 9),
+        (4, 9, 5),
+        (2, 4, 11),
+        (6, 2, 10),
+        (8, 6, 7),
+        (9, 8, 1),
+    ]
+    vertices = [np.array(v) / np.linalg.norm(v) for v in vertices]
+    return vertices, faces
+
+
+def refine(vertices: list, faces: list) -> list:
+    """Split each triangle into four, adding its edges' midpoints, on the
+    sphere, to ``vertices``."""
+    midpoints = {}
+
+    def midpoint(a, b):
+        key = (min(a, b), max(a, b))
+        if key not in midpoints:
+            m = vertices[a] + vertices[b]
+            vertices.append(m / np.linalg.norm(m))
+            midpoints[key] = len(vertices) - 1
+        return midpoints[key]
+
+    finer = []
+    for a, b, c in faces:
+        ab, bc, ca = midpoint(a, b), midpoint(b, c), midpoint(c, a)
+        finer += [(a, ab, ca), (b, bc, ab), (c, ca, bc), (ab, bc, ca)]
+    return finer
+
+
+def multimesh(refinements: int):
+    """The vertices of an icosahedron refined ``refinements`` times, the
+    edges of every level of refinement in both directions, and the length
+    of the finest level's longest edge."""
+    vertices, faces = icosahedron()
+    levels = [faces]
+    for _ in range(refinements):
+        levels.append(refine(vertices, levels[-1]))
+    edges = {
+        (u, v)
+        for level in levels
+        for face in level
+        for u, v in zip(face, face[1:] + face[:1])
+    }
+    edges |= {(v, u) for u, v in edges}
+    vertices = np.array(vertices)
+    finest = max(
+        np.linalg.norm(vertices[u] - vertices[v])
+        for face in levels[-1]
+        for u, v in zip(face, face[1:] + face[:1])
+    )
+    return vertices, np.array(sorted(edges)).T, finest
+
+
+def xyz(latitude, longitude):
+    lat, lon = np.radians(latitude), np.radians(longitude)
+    return np.stack(
+        [np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)], -1
+    )
+
+
+def graph_table(senders, receivers, points_from, points_to) -> xr.Dataset:
+    """A graph as a Dataset: the table ``edges(edge, sender, receiver)``
+    and the table ``features(edge, f, val)``."""
+    offset = points_to[receivers] - points_from[senders]
+    features = np.concatenate(
+        [offset, np.linalg.norm(offset, axis=1, keepdims=True)], axis=1
+    )
+    return xr.Dataset(
+        {
+            "sender": ("edge", senders),
+            "receiver": ("edge", receivers),
+            "val": (("edge", "f"), features),
+        },
+        coords={"edge": np.arange(len(senders)), "f": np.arange(EDGE_FEATURES)},
+    )
+
+
+GRAPH_TABLES = {("edge",): "edges", ("edge", "f"): "features"}
+
+
+def geometry(era5: xr.Dataset, refinements: int) -> dict[str, xr.Dataset]:
+    """The tables of the model's fixed structure, as Datasets."""
+    lat, lon = xr.broadcast(era5.latitude, era5.longitude)
+    lat, lon = lat.values.ravel(), lon.values.ravel()
+    cell_xyz = xyz(lat, lon)
+    node_xyz, (senders, receivers), finest = multimesh(refinements)
+    node_lat = np.degrees(np.arcsin(node_xyz[:, 2]))
+    node_lon = np.degrees(np.arctan2(node_xyz[:, 1], node_xyz[:, 0])) % 360
+
+    # Grid to mesh: each cell to the mesh nodes within 0.6 of the finest
+    # mesh's longest edge, GraphCast's radius.
+    distance = np.linalg.norm(cell_xyz[:, None] - node_xyz[None], axis=-1)
+    g2m_cell, g2m_node = np.nonzero(distance < 0.6 * finest)
+    # Mesh to grid: each cell from its 3 nearest mesh nodes.
+    nearest = np.argsort(distance, axis=1)[:, :3]
+    m2g_cell = np.repeat(np.arange(len(lat)), 3)
+    m2g_node = nearest.ravel()
+
+    return {
+        "cells": xr.Dataset(
+            {"latitude": ("cell", lat), "longitude": ("cell", lon)},
+            coords={"cell": np.arange(len(lat))},
+        ),
+        "mesh": xr.Dataset(
+            {"latitude": ("node", node_lat), "longitude": ("node", node_lon)},
+            coords={"node": np.arange(len(node_xyz))},
+        ),
+        "g2m": graph_table(g2m_cell, g2m_node, cell_xyz, node_xyz),
+        "multimesh": graph_table(senders, receivers, node_xyz, node_xyz),
+        "m2g": graph_table(m2g_node, m2g_cell, node_xyz, cell_xyz),
+    }
+
+
+# --------------------------------------------------------------------------
+# Data and parameters.
+# --------------------------------------------------------------------------
+
+
+def open_era5() -> xr.Dataset:
+    """WeatherBench 2's ERA5, lazily; a synthetic stand-in offline."""
+    try:
+        return xr.open_zarr(
+            ERA5, chunks=None, storage_options={"token": "anon"}
+        )
+    except Exception as e:  # noqa: BLE001 - any failure to reach the store
+        print(f"  ERA5 unavailable ({type(e).__name__}); using synthetic data")
+    lon = np.arange(64) * 5.625
+    lat = -87.1875 + np.arange(32) * 5.625
+    rng = np.random.default_rng(3)
+    surface = rng.normal(size=(len(TIMES), 64, 32))
+    atmosphere = rng.normal(size=(len(TIMES), 2, 64, 32))
+    data: dict[str, tuple[tuple[str, ...], np.ndarray]] = {
+        name: (
+            ("time", "longitude", "latitude"),
+            surface + i,
+        )
+        for i, (table, name, _) in enumerate(VARIABLES)
+        if table == "surface"
+    }
+    data["geopotential"] = (
+        ("time", "level", "longitude", "latitude"),
+        atmosphere,
+    )
+    data["temperature"] = (
+        ("time", "level", "longitude", "latitude"),
+        -atmosphere,
+    )
+    return xr.Dataset(
+        data,
+        coords={
+            "time": TIMES,
+            "level": [500, 850],
+            "longitude": lon,
+            "latitude": lat,
+        },
+    )
+
+
+def era5_arrays(era5: xr.Dataset, time) -> np.ndarray:
+    """Each variable at ``time`` as a column, one row per cell, read with
+    xarray: the reference's input, independent of the SQL's."""
+    columns = []
+    for _, name, level in VARIABLES:
+        field = era5[name].sel(time=time)
+        if level:
+            field = field.sel(level=level)
+        columns.append(field.transpose("latitude", "longitude").values.ravel())
+    return np.stack(columns, axis=1)
+
+
+def init_params(rounds: int, rng) -> dict[str, np.ndarray]:
+    """Each network's weight matrices, by name, with fan-in scaling."""
+    shapes = {"output": (HIDDEN, len(VARIABLES))}
+
+    def two_layer(net, inputs):
+        shapes[f"{net}.0"] = (inputs, HIDDEN)
+        shapes[f"{net}.1"] = (HIDDEN, HIDDEN)
+
+    two_layer("embed.grid", INPUTS)
+    two_layer("embed.mesh", 3)
+    for graph in ("g2m", "multimesh", "m2g"):
+        two_layer(f"embed.{graph}", EDGE_FEATURES)
+    for net in ["encoder", "decoder"] + [
+        f"processor.{r}" for r in range(rounds)
+    ]:
+        two_layer(f"{net}.edge", 3 * HIDDEN)
+        two_layer(f"{net}.node", 2 * HIDDEN)
+    return {
+        net: rng.normal(0, 1 / np.sqrt(shape[0]), shape)
+        for net, shape in shapes.items()
+    }
+
+
+def params_table(params: dict[str, np.ndarray]) -> xr.Dataset:
+    """Every weight as a row of ``params(weight, net, f, g, w)``."""
+    rows = [
+        (net, f, g, w)
+        for net, matrix in params.items()
+        for (f, g), w in np.ndenumerate(matrix)
+    ]
+    net, f, g, w = map(np.array, zip(*rows))
+    return xr.Dataset(
+        {
+            "net": ("weight", net),
+            "f": ("weight", f),
+            "g": ("weight", g),
+            "w": ("weight", w),
+        },
+        coords={"weight": np.arange(len(rows))},
+    )
+
+
+def main():
+    parser = cli(__doc__.splitlines()[0], sizes=[2], size_name="refinements")
+    parser.add_argument(
+        "--rounds", type=int, default=2, help="processor rounds"
+    )
+    args = parser.parse_args()
+    report = Report("graphcast", "refinements", args.reps)
+    rounds = args.rounds
+    era5 = open_era5()
+
+    # The normalization: each variable's mean and spread at the inputs.
+    now, before, later = (era5_arrays(era5, t) for t in TIMES[[1, 0, 2]])
+    mean = np.concatenate([before, now]).mean(axis=0)
+    std = np.concatenate([before, now]).std(axis=0)
+    scale = xr.Dataset(
+        {"mean": ("var", mean), "std": ("var", std)},
+        coords={"var": np.arange(len(VARIABLES))},
+    )
+    params = init_params(rounds, np.random.default_rng(0))
+
+    for refinements in args.sizes:
+        tables = geometry(era5, refinements)
+        cells = tables["cells"]
+        print(
+            f"  mesh refined {refinements}×: {tables['mesh'].sizes['node']} "
+            f"nodes, {tables['multimesh'].sizes['edge']} mesh edges, "
+            f"{tables['g2m'].sizes['edge']} grid-to-mesh, "
+            f"{tables['m2g'].sizes['edge']} mesh-to-grid; "
+            f"{len(params_table(params).weight)} weights"
+        )
+
+        # The reference, in JAX, from arrays read with xarray.
+        graph = {
+            name: {
+                "edges": (
+                    tables[name].sender.values,
+                    tables[name].receiver.values,
+                ),
+                "features": tables[name].val.values,
+            }
+            for name in ("g2m", "multimesh", "m2g")
+        }
+        graph["mesh_in"] = np.stack(
+            [
+                np.cos(np.radians(tables["mesh"].latitude.values)),
+                np.sin(np.radians(tables["mesh"].longitude.values)),
+                np.cos(np.radians(tables["mesh"].longitude.values)),
+            ],
+            axis=1,
+        )
+        position = np.stack(
+            [
+                np.cos(np.radians(cells.latitude.values)),
+                np.sin(np.radians(cells.longitude.values)),
+                np.cos(np.radians(cells.longitude.values)),
+            ],
+            axis=1,
+        )
+        grid_in = np.concatenate(
+            [(before - mean) / std, (now - mean) / std, position], axis=1
+        )
+        target = (later - now) / std
+        want = jax_forecast(params, graph, grid_in, rounds)
+
+        want_grad = jax.grad(jax_error)(params, graph, grid_in, target, rounds)
+        want_grad = np.concatenate([want_grad[net].ravel() for net in params])
+
+        for host in hosts(args.hosts):
+            host.load("era5", era5, table_names=ERA5_TABLES)
+            for name, table in tables.items():
+                options = (
+                    {"table_names": GRAPH_TABLES}
+                    if "edge" in table.dims
+                    else {}
+                )
+                host.load(name, table, **options)
+            host.load("scale", scale)
+            host.load("params", params_table(params))
+
+            statement = forecast_statement(rounds)
+            rows = report.time(
+                "forecast, global", host, refinements, host.query, statement
+            )
+            got = rows["change"].to_numpy().reshape(len(cells.cell), -1)
+            # The rows are ordered by latitude then longitude, as cells are.
+            report.check(assert_close("forecast", got, want, rtol=1e-9))
+            report.note("operators", host.operators(statement))
+
+            statement = forecast_statement(rounds, REGION)
+            rows = report.time(
+                "forecast, regional", host, refinements, host.query, statement
+            )
+            in_region = np.logical_and.reduce(
+                [
+                    (cells[c] >= lo) & (cells[c] <= hi)
+                    for c, (lo, hi) in BOUNDS.items()
+                ]
+            )
+            regional = rows["change"].to_numpy().reshape(-1, len(VARIABLES))
+            report.check(assert_close("regional", regional, want[in_region]))
+            report.note("operators", host.operators(statement))
+
+            step = "gradient of the error"
+            if not host.differentiates:
+                report.skip(step, host, refinements, "no ddx backend")
+                continue
+            try:
+                rows = report.time(
+                    step,
+                    host,
+                    refinements,
+                    host.query,
+                    gradient_statement(rounds),
+                )
+            except DdxError as e:
+                # Recorded rather than raised: a gap in ddx, which this
+                # row tracks until ddx closes it.
+                reason = str(e).removeprefix("not supported by ddx-ad yet: ")
+                report.skip(step, host, refinements, f"ddx: {reason}")
+                continue
+            report.check(
+                assert_close("gradient", rows["w"].to_numpy(), want_grad)
+            )
+
+        print("\n  receptive field of the regional forecast:")
+        for stage, n in receptive_field(host, rounds):
+            print(f"    {stage:<34} {n:>6}")
+
+    report.finish(args.csv)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md
new file mode 100644
index 00000000..40a92a6a
--- /dev/null
+++ b/benchmarks/ml/README.md
@@ -0,0 +1,188 @@
+# Machine learning in SQL
+
+These scripts write three machine-learning models as SQL over tables that
+[xarray-sql](../../README.md) registers from xarray Datasets, and time the
+same SQL on several engines. Two related projects shape them:
+
+- [ddx](https://github.com/xqlsystems/ddx) differentiates SQL queries: it
+  turns the gradient of a query into more SQL, which trains the models.
+- [einfold](https://github.com/xqlsystems/einfold) is a query optimizer
+  that rewrites SQL plans into array algorithms. These models are its
+  three [target demos](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md),
+  and this SQL is the baseline its plans must beat.
+
+Each script checks its results against the same model in
+[JAX](https://github.com/jax-ml/jax), Google's library for differentiable
+array programs.
+
+| # | Script | Model | What it measures |
+|---|---|---|---|
+| 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, as a view (the model) in einfold's window and `GROUP BY` spellings, and a loss that reads it | the naive plan, which stores every score, for [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan to beat; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention, the Nadaraya–Watson kernel average |
+| 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, as a view, and a cross-entropy loss that reads it, on Tiny Shakespeare | the loss on each host; training steps with SGD or, as nanoGPT trains, AdamW, every parameter table checked after every step |
+| 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature [GraphCast](https://arxiv.org/abs/2212.12794), DeepMind's graph neural network for weather forecasting, over [ERA5](https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5) reanalysis in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region |
+
+The SQL favors legibility over speed. Each model reads like its definition,
+and the slow parts are left slow on purpose: they are what ddx, xarray-sql
+and einfold need to make fast.
+
+Every result is checked before it is timed into the table. Each script
+computes the same model in JAX from the same inputs and parameters, and a
+mismatch fails the run.
+
+## Running
+
+```shell
+uv run benchmarks/ml/01_attention.py
+uv run benchmarks/ml/02_nanogpt.py --steps 5
+uv run benchmarks/ml/03_graphcast.py --rounds 1
+```
+
+Each script carries [inline `uv` metadata](https://docs.astral.sh/uv/guides/scripts/)
+and runs against the xarray-sql in this checkout, with the `ml` extra. The
+options are the same in all three:
+
+```shell
+--hosts datafusion duckdb   # the engines to run on (default: all)
+--sizes 64 256              # problem sizes: tokens, batch size, or mesh refinements
+--reps 3                    # timed repetitions; the median is reported
+--csv results.csv           # also write the results table
+```
+
+Demo 2 adds `--steps`, the number of training steps, and `--optimizer`:
+`sgd` (the default), or `adamw` to train as nanoGPT does. Demo 3 adds
+`--rounds`, the number of processor rounds.
+
+The output ends with one table: a row per step, host and size, with its
+median time, its error against JAX, and the number of operators in the
+host's physical plan.
+
+## The `ml` extra
+
+```shell
+pip install "xarray-sql[ml]"
+```
+
+The extra installs [ddxdb](https://pypi.org/project/ddxdb/). With it, an
+`xarray_sql.Context` understands `grad`, `vjp` and `jvp` in a select list,
+and `ddx_stop_gradient` as a function. A gradient-descent step on a
+parameter table is one statement:
+
+```sql
+WITH pred AS (
+  SELECT x.sample, SUM(x.value * w.val) AS y
+  FROM x JOIN w ON x.feature = w.feature GROUP BY x.sample),
+loss AS (
+  SELECT AVG(power(pred.y - t.value, 2)) AS l
+  FROM pred JOIN t ON pred.sample = t.sample)
+SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w
+```
+
+`xarray_sql.ad.sql_all(ctx, statements)` runs several such statements for
+one backward pass, which is how Demo 2 updates its eleven parameter tables.
+
+## Hosts
+
+einfold calls the engine that runs a plan its *host*.
+[`_hosts.py`](_hosts.py) defines each one in a few lines: how it registers
+an `xr.Dataset` (through xarray-sql), runs SQL, and counts the operators
+in its plan.
+
+| Host | Forward pass | Gradients |
+|---|---|---|
+| `datafusion` | `xarray_sql.Context` | ddx |
+| `duckdb` | DuckDB, through `xql.register` | not yet: ddx runs its programs on DataFusion only |
+
+A host that cannot differentiate still runs every forward query, and its
+gradient rows say `no ddx backend`. To add a host, such as DataFusion with
+einfold's optimizer or DuckDB with a GPU extension, write one class in
+`_hosts.py` and add it to `HOSTS`.
+
+## Findings
+
+Measured on 2026-10-17 on an Intel Core i7-8700 (6 cores, 12 threads)
+with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each.
+The GPT and attention demos were timed again on 2026-10-21 with
+DataFusion 55.0.0 and ddxdb 0.4.0, after xarray-sql began registering
+small in-memory Datasets as DataFusion in-memory tables.
+
+**ddx's gradients agree with JAX.** The gradients of attention with
+respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the
+largest value at every size. After each of five SGD steps, all eleven of
+the GPT's parameter tables match JAX's to 1e-15, and the loss falls from
+4.18 to 3.84 (each step on a different batch). With AdamW, the weights
+and both of Adam's moments match a JAX transcription of
+`torch.optim.AdamW` to 6e-15 after each of five steps, and that
+transcription matches PyTorch exactly.
+
+**DataFusion computes a CTE once for every place that reads it.** A
+model reads its activations more than once: a layer norm reads its input
+three times, a residual connection twice, a message-passing round reads
+the mesh as sender, receiver and residual. DataFusion inlines each read,
+so the plan grows with the model's depth; DuckDB computes a CTE that is
+read more than once a single time.
+
+| Query | DataFusion operators | DuckDB operators | DataFusion | DuckDB |
+|---|---:|---:|---:|---:|
+| GPT loss, 1 layer, batch 4 | 13,756 | 126 | 3.0 s | 0.057 s |
+| GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s |
+| GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s |
+
+The GraphCast rows are from 2026-10-17, before in-memory tables.
+
+Setting up each lazy read of a Dataset takes xarray-sql about 1.5 ms, so
+the GPT's 2,989 reads in one forward pass took 6.7 s of its 9.4 s. A small
+in-memory Dataset is now an in-memory table, which takes a quarter of a
+millisecond per read. A GPT training step, forward and backward through
+ddx, takes 0.75 s on DataFusion with SGD and 0.85 s with AdamW, a quarter
+of its forward query alone (1.2 s and 1.55 s with ddxdb 0.3.3 on
+DataFusion 54). ddx builds the gradient program on the first step, which
+takes 1.2 s, and runs it again on later ones. It plans each CTE the loss
+reads twice on its own, and stores the result of each step of its program
+as a table, so it computes each shared subexpression once. Doing the same
+for forward queries is the first thing xarray-sql's DataFusion path needs
+for these models, and einfold should plan them as graphs rather than
+trees.
+
+**Attention stores every score.** Causal attention's forward pass at 512
+tokens takes 0.28 s on DataFusion in the window spelling and 0.37 s in the
+`GROUP BY` spelling; DuckDB takes 0.48 s and 0.27 s. Both grow with the
+square of the tokens. The gradients take 2.5 to 3.5 times the forward
+pass: 0.92 s and 0.93 s on DataFusion at 512 tokens, against 1.1 s and 1.3
+s with ddxdb 0.3.3 on DataFusion 54, and 2.8 s and 2.7 s with ddxdb 0.3.2,
+before ddx differentiated the masked product as `jax.grad` does.
+
+**A regional forecast costs what a global one does.** For the contiguous
+United States, the forecast covers 55 of 2,048 grid cells. Walking the
+graph back from them, the decoder needs 15 mesh nodes, the processor's
+two rounds need 36 and then 82, and the encoder reads 1,103 grid cells
+(54%); with one round, 576 (28%). Neither host uses any of this: the
+regional forecast takes as long as the global one on both. Pushing the
+region's filter back through the graph is einfold's Demo 3.
+
+**What ddx cannot differentiate yet.** Each of these raised a typed
+error naming the construct, never a wrong number:
+
+- A `GROUP BY` key that depends on a differentiated column. einfold's
+  window spelling groups by the normalizer `z`; the demos divide inside
+  the sum instead, `SUM(w / z * v)`, which einfold also lists.
+- `UNION ALL` between a parameter table and the loss. GraphCast
+  concatenates features with `UNION ALL`, so its gradient row records the
+  error. It passes once ddx differentiates a union.
+- Any engine but DataFusion: ddx's DuckDB support is planned.
+
+**xarray-sql reads NaN as missing.** A NaN in a Dataset becomes SQL
+`NULL`, which `SUM` skips, so Demo 1 checks masks with an infinity: a
+plan that multiplies a masked pair's zero weight by it gives NaN where
+SQL gives a number.
+
+## Next
+
+- Part B of the NanoGPT demo: the four tricks behind a record of the
+  [NanoGPT speedrun](https://github.com/KellerJordan/modded-nanogpt), a
+  public race to train GPT-2 fastest, as a small diff to `02_nanogpt.py`.
+  Its sparse optimizer updates build on `--optimizer adamw`.
+- Peak memory per host, from a fresh process per measurement.
+- An einfold host, once einfold produces plans.
+
+[`_bench.py`](_bench.py) holds the shared command line, timing, checks and
+results table.
diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py
new file mode 100644
index 00000000..eed1f918
--- /dev/null
+++ b/benchmarks/ml/_bench.py
@@ -0,0 +1,151 @@
+"""Timing, checking and reporting for the ML demos.
+
+Each demo is a script with the same command line::
+
+    uv run benchmarks/ml/01_attention.py --hosts datafusion duckdb \\
+        --sizes 64 256 --reps 3 --csv attention.csv
+
+It runs every step on every chosen host (an engine that runs SQL; see
+``_hosts.py``) and size, checks each result against the same model in
+JAX, Google's library for differentiable array programs, and ends with
+one table: a row per step, host
+and size, with the median time and whether the result matched. Checks
+come first: a wrong answer fails the run, however fast it was.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import statistics
+import sys
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+
+import numpy as np
+from _hosts import HOSTS, Host
+
+
+def cli(description: str, sizes: Sequence[int], size_name: str = "n"):
+    """The shared command line, for a demo to extend and parse."""
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument(
+        "--hosts",
+        nargs="+",
+        default=list(HOSTS),
+        choices=list(HOSTS),
+        help="engines to run on (default: all)",
+    )
+    parser.add_argument(
+        "--sizes",
+        nargs="+",
+        type=int,
+        default=list(sizes),
+        help=f"problem sizes, as {size_name} (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--reps",
+        type=int,
+        default=1,
+        help="timed repetitions per step; the median is reported",
+    )
+    parser.add_argument("--csv", help="also write the results table here")
+    return parser
+
+
+def hosts(names: Sequence[str]) -> list[Host]:
+    """A fresh instance of each named host."""
+    return [HOSTS[name]() for name in names]
+
+
+def assert_close(name: str, got, want, *, rtol=1e-9, atol=1e-9) -> float:
+    """Assert two arrays match, NaN for NaN. Returns the largest error,
+    relative to the largest value."""
+    got, want = np.asarray(got), np.asarray(want)
+    np.testing.assert_allclose(
+        got, want, rtol=rtol, atol=atol, equal_nan=True, err_msg=name
+    )
+    finite = np.isfinite(want)
+    scale = np.max(np.abs(want[finite]), initial=0.0) or 1.0
+    return float(np.max(np.abs(got - want)[finite], initial=0.0) / scale)
+
+
+@dataclass
+class Report:
+    """The demo's results: one row per step, host and size."""
+
+    demo: str
+    size_name: str
+    reps: int = 1
+    rows: list[dict] = field(default_factory=list)
+
+    def time(self, step: str, host: Host, size: int, run: Callable, *args):
+        """Run ``run(*args)`` ``reps`` times, record the median, and return
+        the last result."""
+        times = []
+        for _ in range(self.reps):
+            start = time.perf_counter()
+            result = run(*args)
+            times.append(time.perf_counter() - start)
+        self.rows.append(
+            {
+                "demo": self.demo,
+                "step": step,
+                "host": host.name,
+                self.size_name: size,
+                "seconds": statistics.median(times),
+                "check": "",
+            }
+        )
+        print(
+            f"  {step:<28} {host.name:<11} {self.size_name}={size:<6} "
+            f"{statistics.median(times):8.3f}s",
+            flush=True,
+        )
+        return result
+
+    def check(self, error: float) -> None:
+        """Mark the last row's result as checked, with its error (from
+        :func:`assert_close`)."""
+        self.rows[-1]["check"] = f"ok ({error:.0e})"
+
+    def skip(self, step: str, host: Host, size: int, why: str) -> None:
+        """Record a step the host cannot run, and why."""
+        self.rows.append(
+            {
+                "demo": self.demo,
+                "step": step,
+                "host": host.name,
+                self.size_name: size,
+                "seconds": None,
+                "check": why,
+            }
+        )
+
+    def note(self, key: str, value) -> None:
+        """Attach a measurement besides time to the last row."""
+        self.rows[-1][key] = value
+
+    def finish(self, csv_path: str | None = None) -> None:
+        """Print the table, and write it as CSV if asked."""
+        columns = list(dict.fromkeys(k for row in self.rows for k in row))
+        print(f"\n| {' | '.join(columns)} |")
+        print(f"|{'|'.join('---' for _ in columns)}|")
+        for row in self.rows:
+            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)
+            print(f"\nwrote {csv_path}", file=sys.stderr)
+
+
+def _cell(value) -> str:
+    if value is None:
+        return "n/a"
+    if isinstance(value, float):
+        return f"{value:.3f}"
+    return str(value)
diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py
new file mode 100644
index 00000000..c52fb71b
--- /dev/null
+++ b/benchmarks/ml/_hosts.py
@@ -0,0 +1,120 @@
+"""The SQL engines ("hosts") the ML demos run on.
+
+einfold (https://github.com/xqlsystems/einfold), a query optimizer that
+rewrites SQL plans into array algorithms, calls the engine that runs a
+plan its *host*. Every demo writes
+its model once, as SQL, and runs that same text on each host, so a
+difference in time is a difference between engines, never between
+queries. A host does five things:
+
+* ``load(name, data, **options)`` registers an ``xr.Dataset`` as the
+  table ``name`` through xarray-sql, replacing any table of that name
+  (``options`` such as ``table_names`` go to xarray-sql);
+* ``execute(sql)`` runs a statement for its effect, such as
+  ``CREATE VIEW``;
+* ``query(sql)`` runs a statement and returns its rows as Arrow;
+* ``query_all(sqls)`` runs several statements, and, on a host that
+  differentiates, shares the backward pass between those that take
+  ``grad`` of the same loss;
+* ``operators(sql)`` counts the operators in the physical plan it would
+  run, which shows, for one, whether a CTE read twice is computed twice.
+
+Forward queries run everywhere. Gradients (``grad``, ``vjp``, ``jvp``)
+need ddx (https://github.com/xqlsystems/ddx), which differentiates SQL
+queries and today runs its programs on DataFusion only; a host that
+cannot differentiate says so with ``differentiates = False``, and the
+demos report its gradient steps as not run rather than skip it silently.
+
+Adding a host (einfold's optimizer on DataFusion, DuckDB with a GPU
+extension, or DuckDB running ddx's programs once ddx supports it) is one
+class here and one entry in ``HOSTS``.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+
+import duckdb
+import pyarrow as pa
+import xarray as xr
+
+import xarray_sql as xql
+
+
+class Host:
+    """An engine that runs the demos' SQL. Subclasses fill in the rest."""
+
+    name = ""
+    differentiates = False
+
+    def load(self, name: str, data: xr.Dataset, **options) -> None:
+        raise NotImplementedError
+
+    def execute(self, sql: str) -> None:
+        raise NotImplementedError
+
+    def query(self, sql: str) -> pa.Table:
+        raise NotImplementedError
+
+    def query_all(self, sqls: Sequence[str]) -> list[pa.Table]:
+        return [self.query(sql) for sql in sqls]
+
+    def operators(self, sql: str) -> int:
+        raise NotImplementedError
+
+
+class DataFusion(Host):
+    """xarray-sql's ``Context``, with ddx's gradients (the ``ml`` extra)."""
+
+    name = "datafusion"
+    differentiates = True
+
+    def __init__(self):
+        self.ctx = xql.Context()
+
+    def load(self, name, data, **options):
+        self.ctx.deregister_table(name)
+        self.ctx.from_dataset(name, data, **options)
+
+    def execute(self, sql):
+        self.ctx.sql(sql).collect()
+
+    def query(self, sql):
+        return self.ctx.sql(sql).to_arrow_table()
+
+    def query_all(self, sqls):
+        return [df.to_arrow_table() for df in xql.ad.sql_all(self.ctx, sqls)]
+
+    def operators(self, sql):
+        plan = self.ctx.sql(sql).execution_plan().display_indent()
+        return len(plan.splitlines())
+
+
+class DuckDB(Host):
+    """DuckDB, reading the same Datasets through xarray-sql's adapter."""
+
+    name = "duckdb"
+
+    def __init__(self):
+        self.con = duckdb.connect()
+
+    def load(self, name, data, **options):
+        xql.register(self.con, name, data, **options)
+
+    def execute(self, sql):
+        self.con.execute(sql)
+
+    def query(self, sql):
+        return self.con.sql(sql).fetch_arrow_table()
+
+    def operators(self, sql):
+        (_, plan), *_ = self.con.sql(f"EXPLAIN (FORMAT JSON) {sql}").fetchall()
+        return sum(_count(node) for node in json.loads(plan))
+
+
+def _count(node: dict) -> int:
+    return 1 + sum(_count(child) for child in node.get("children", []))
+
+
+HOSTS = {host.name: host for host in (DataFusion, DuckDB)}
diff --git a/pyproject.toml b/pyproject.toml
index 8d80eae7..77bb1b0e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -58,11 +58,17 @@ polars = [
 geo = [
   "pyproj",
 ]
+# Gradients in SQL: grad(), vjp() and jvp() from ddx. ddxdb is pre-1.0
+# and breaks its API between minor versions. It reads and writes the
+# Substrait of DataFusion 55, so bump it together with that pin.
+ml = [
+  "ddxdb>=0.4.0,<0.5; sys_platform != 'emscripten'",
+]
 test = [
   "adbc-driver-flightsql>=1.12",
   "adbc-driver-sqlite>=1.12",
   "cftime",
-  "xarray-sql[adbc,dask,duckdb,polars,geo]",
+  "xarray-sql[adbc,dask,duckdb,polars,geo,ml]",
   "pytest",
   "xarray[io]",
   "gcsfs",
diff --git a/tests/test_ad.py b/tests/test_ad.py
new file mode 100644
index 00000000..b395a247
--- /dev/null
+++ b/tests/test_ad.py
@@ -0,0 +1,99 @@
+"""Tests for the ddx gradient extension (`xarray_sql.ad`)."""
+
+import numpy as np
+import pytest
+import xarray as xr
+
+pytest.importorskip("ddxdb")
+
+from ddxdb import DdxError
+
+from xarray_sql import Context, ad
+
+# Least squares: pred = x @ w, loss = sum((pred - y)^2).
+LOSS = """
+pred AS (
+  SELECT x.sample, SUM(x.x * w.w) AS p
+  FROM data x JOIN weights w ON x.feature = w.feature
+  GROUP BY x.sample),
+loss AS (
+  SELECT SUM(power(pred.p - y.y, 2)) AS l
+  FROM pred JOIN targets y ON pred.sample = y.sample)
+"""
+
+
+@pytest.fixture
+def problem():
+    rng = np.random.default_rng(0)
+    x = rng.normal(size=(6, 3))
+    w = rng.normal(size=3)
+    y = rng.normal(size=6)
+    sample, feature = np.arange(6), np.arange(3)
+    ctx = Context()
+    ctx.from_dataset(
+        "data",
+        xr.Dataset(
+            {"x": (("sample", "feature"), x)},
+            coords={"sample": sample, "feature": feature},
+        ),
+    )
+    ctx.from_dataset(
+        "weights",
+        xr.Dataset({"w": ("feature", w)}, coords={"feature": feature}),
+    )
+    ctx.from_dataset(
+        "targets", xr.Dataset({"y": ("sample", y)}, coords={"sample": sample})
+    )
+    return ctx, x, w, y
+
+
+def test_grad_matches_closed_form(problem):
+    ctx, x, w, y = problem
+    got = ctx.sql(
+        f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature"
+    ).to_dataset()
+    np.testing.assert_allclose(got["w"].values, 2 * x.T @ (x @ w - y))
+
+
+def test_gradient_step_reads_like_sql(problem):
+    ctx, x, w, y = problem
+    got = ctx.sql(
+        f"WITH {LOSS} "
+        "SELECT feature, w - 0.01 * grad(loss, w) AS w FROM weights "
+        "ORDER BY feature"
+    ).to_pandas()
+    np.testing.assert_allclose(got["w"], w - 0.01 * 2 * x.T @ (x @ w - y))
+
+
+def test_sql_all_shares_one_loss(problem):
+    ctx, x, w, y = problem
+    grad = 2 * x.T @ (x @ w - y)
+    select = "SELECT feature, {} AS w FROM weights ORDER BY feature"
+    gradient, step = ad.sql_all(
+        ctx,
+        [
+            f"WITH {LOSS} " + select.format("grad(loss, w)"),
+            f"WITH {LOSS} " + select.format("w - 0.01 * grad(loss, w)"),
+        ],
+    )
+    np.testing.assert_allclose(gradient.to_dataset()["w"].values, grad)
+    np.testing.assert_allclose(step.to_dataset()["w"].values, w - 0.01 * grad)
+
+
+def test_stop_gradient_blocks_the_derivative(problem):
+    ctx, _, w, _ = problem
+    got = ctx.sql(
+        "WITH loss AS (SELECT SUM(w * ddx_stop_gradient(w)) AS l "
+        "FROM weights) "
+        "SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature"
+    ).to_dataset()
+    np.testing.assert_allclose(got["w"].values, w)
+
+
+def test_unsupported_expression_raises(problem):
+    ctx, *_ = problem
+    with pytest.raises(DdxError):
+        ctx.sql(
+            "WITH loss AS (SELECT SUM(w) AS l FROM weights) "
+            "SELECT feature, grad(loss, feature) AS g FROM weights"
+        ).collect()
diff --git a/uv.lock b/uv.lock
index f98b1b8c..879bc3e7 100644
--- a/uv.lock
+++ b/uv.lock
@@ -520,6 +520,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/4b/d0/bdc7ea7072c3925f3977e13e1e754e145fa42eb40bb5c1abff8028e0f791/datafusion-55.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b7191ff60defee8d9d9e5ff6aeefdd961bc200897cd60440b02d2a0f316902ef", size = 45733022, upload-time = "2026-10-12T16:20:42.607Z" },
 ]
 
+[[package]]
+name = "ddxdb"
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/92/9b/e545bc8537ec63c8c41d0642e52798956eda3393feaeccb0406a3a5950c7/ddxdb-0.4.0.tar.gz", hash = "sha256:0c0d4a14b82af02d6b83b9ed243014c64401a167f123c8a84d9fc19429d5f1ed", size = 291824, upload-time = "2026-10-21T09:04:11.42Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c4/7e/acbeeca83355880e4404a7b15d765ed82d85efa4a79952941c9401931885/ddxdb-0.4.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:be538a452d2c45577fb060a2be3134631ea68658caa9c3a2a854b5b13a0f21d9", size = 6725712, upload-time = "2026-10-21T09:04:00.956Z" },
+    { url = "https://files.pythonhosted.org/packages/23/b9/0fc5bbb27bd323533089c6372adc80f79875679f7741a9d4e738eabff1b2/ddxdb-0.4.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:8c4a752f27db4219ad0bd8d6b301e88bea192bc3753b54bc909dd81bbf562781", size = 6644600, upload-time = "2026-10-21T09:04:03.238Z" },
+    { url = "https://files.pythonhosted.org/packages/ae/8a/b76f44799346a5d6b9f88c1dde5d2aae19934ce31b6c992b27629eee8709/ddxdb-0.4.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0636c7d4d52e64e0d3bc71dcf1bb3a15d1e1e494c553e6b184da192df147be69", size = 6739521, upload-time = "2026-10-21T09:04:05.205Z" },
+    { url = "https://files.pythonhosted.org/packages/f5/cc/2b80f26e102beb2c655dc2b60bf2e427b0d9e0e588c723343c8ea300dac4/ddxdb-0.4.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b2f745b7cd057f0a7fc550f0f5984a836683a6e3252a677fc18e6f6045dbb80d", size = 7179153, upload-time = "2026-10-21T09:04:07.276Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/dc/0231a760cbea1a86a8c71d0526ec9ff45d7850573ea73aacddcf3d9916b6/ddxdb-0.4.0-cp310-abi3-win_amd64.whl", hash = "sha256:f38c950fdc8ec1a4c347d80715c6b3c04499922f2e71dac004ef51b2adc71101", size = 7133649, upload-time = "2026-10-21T09:04:09.43Z" },
+]
+
 [[package]]
 name = "decorator"
 version = "5.2.1"
@@ -2865,6 +2878,9 @@ geo = [
     { name = "pyproj", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "pyproj", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
 ]
+ml = [
+    { name = "ddxdb", marker = "sys_platform != 'emscripten'" },
+]
 polars = [
     { name = "polars" },
 ]
@@ -2874,6 +2890,7 @@ test = [
     { name = "adbc-driver-sqlite" },
     { name = "cftime" },
     { name = "dask" },
+    { name = "ddxdb", marker = "sys_platform != 'emscripten'" },
     { name = "duckdb" },
     { name = "gcsfs" },
     { name = "polars" },
@@ -2900,6 +2917,7 @@ requires-dist = [
     { name = "cftime", marker = "extra == 'test'" },
     { name = "dask", marker = "extra == 'dask'", specifier = ">=2024.8.0" },
     { name = "datafusion", marker = "sys_platform != 'emscripten'", specifier = "==55.0.0" },
+    { name = "ddxdb", marker = "sys_platform != 'emscripten' and extra == 'ml'", specifier = ">=0.4.0,<0.5" },
     { name = "duckdb", marker = "extra == 'duckdb'", specifier = ">=1.4.0" },
     { name = "gcsfs", marker = "extra == 'test'" },
     { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'" },
@@ -2912,11 +2930,11 @@ requires-dist = [
     { name = "watchfiles", marker = "extra == 'dev'" },
     { name = "xarray", specifier = ">=2024.7.0" },
     { name = "xarray", extras = ["io"], marker = "extra == 'test'" },
-    { name = "xarray-sql", extras = ["adbc", "dask", "duckdb", "polars", "geo"], marker = "extra == 'test'" },
+    { name = "xarray-sql", extras = ["adbc", "dask", "duckdb", "polars", "geo", "ml"], marker = "extra == 'test'" },
     { name = "xarray-sql", extras = ["docs"], marker = "extra == 'dev'" },
     { name = "zensical", marker = "extra == 'docs'" },
 ]
-provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "polars", "test"]
+provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "ml", "polars", "test"]
 
 [package.metadata.requires-dev]
 dev = [
diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py
new file mode 100644
index 00000000..edf4b09e
--- /dev/null
+++ b/xarray_sql/ad.py
@@ -0,0 +1,133 @@
+"""Gradients of SQL queries: the optional ml extension, built on ddx.
+
+[ddx](https://github.com/xqlsystems/ddx) differentiates a whole SQL query
+the way [JAX](https://github.com/jax-ml/jax), Google's library for
+differentiable array programs, differentiates a function. With it installed
+(``pip install xarray-sql[ml]``), [xarray_sql.Context][]'s ``.sql()``
+understands three calls in a select list:
+
+* ``grad(loss, column)``: at each row, the derivative of the one number
+  the CTE ``loss`` computes with respect to that row's ``column``
+  (``table.column`` when the ``SELECT`` reads several tables).
+* ``vjp(f, column, cotangent)``: the same for any CTE ``f``, pulled back
+  along ``cotangent``, a table or CTE shaped like ``f``'s output.
+* ``jvp(e, table.column, tangent)``: the tangent of the expression ``e``
+  along ``tangent``, pushed forward.
+
+A gradient-descent step on a parameter table is then one statement::
+
+    WITH pred AS (
+      SELECT x.sample, SUM(x.value * w.val) AS y
+      FROM x JOIN w ON x.feature = w.feature GROUP BY x.sample),
+    loss AS (
+      SELECT AVG(power(pred.y - t.value, 2)) AS l
+      FROM pred JOIN t ON pred.sample = t.sample)
+    SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w
+
+The query being differentiated is ordinary SQL, over tables registered
+with ``from_dataset`` or any other way. ddx plans the backward pass in
+[Substrait](https://substrait.io), a portable format for query plans;
+DataFusion runs it, and the statement reads what it computed.
+
+The one function ddx gives a meaning to is ``ddx_stop_gradient(x)``, the
+identity whose derivative is zero (JAX's ``lax.stop_gradient``). It is
+registered on every ``Context``. Its typical use is the stable softmax,
+whose shift by the row maximum cancels and so carries no gradient::
+
+    exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY i)))
+
+A statement that calls none of ``grad``, ``vjp`` or ``jvp`` is planned
+exactly as it would be without ddx. An expression ddx cannot
+differentiate raises a ``ddxdb.DdxError`` subclass naming it, never a
+wrong number.
+
+[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
+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``.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+
+from datafusion import SessionContext
+from ddxdb import ad as _ddx
+
+__all__ = ["STOP_GRADIENT", "calls_ad", "register", "sql", "sql_all"]
+
+STOP_GRADIENT = _ddx.STOP_GRADIENT
+"""The SQL name of the stop-gradient function, ``ddx_stop_gradient``."""
+
+_AD_CALL = re.compile(r"\b(grad|vjp|jvp)\s*\(", re.IGNORECASE)
+
+
+def register(ctx: SessionContext) -> None:
+    """Register ``ddx_stop_gradient`` on ``ctx``.
+
+    Works on any DataFusion ``SessionContext`` (``Context`` registers it
+    automatically when ddxdb is installed). ``grad``, ``vjp`` and ``jvp``
+    are not UDFs: [sql][xarray_sql.ad.sql] rewrites them before DataFusion
+    plans the statement.
+
+    Args:
+        ctx: The DataFusion session context to register the UDF on.
+    """
+    _ddx.register_stop_gradient(ctx)
+
+
+def calls_ad(statement: str) -> bool:
+    """Whether ``statement`` may call ``grad``, ``vjp`` or ``jvp``.
+
+    A cheap textual check, so statements without a call never reach
+    ddx's parser. A false positive (the word inside a string literal,
+    say) is harmless: ddx plans a statement with no call as it is.
+    """
+    return _AD_CALL.search(statement) is not None
+
+
+def sql(ctx: SessionContext, statement: str):
+    """Run ``statement``, computing its ``grad``, ``vjp`` and ``jvp`` calls.
+
+    Returns a DataFusion DataFrame; ``Context.sql`` wraps it as an
+    [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame].
+
+    Args:
+        ctx: Any DataFusion ``SessionContext``.
+        statement: A SQL statement.
+    """
+    return _ddx.sql(ctx, statement)
+
+
+def sql_all(ctx: SessionContext, statements: Sequence[str]) -> list:
+    """Run several statements, sharing the programs their calls need.
+
+    Statements that take ``grad`` of the same loss share one backward
+    pass, so a training step can update each parameter table with its
+    own statement and pay for the gradient once::
+
+        steps = xql.ad.sql_all(ctx, [
+            f"WITH loss AS ({LOSS}) "
+            "SELECT i, val - 0.1 * grad(loss, val) AS val FROM w1",
+            f"WITH loss AS ({LOSS}) "
+            "SELECT j, val - 0.1 * grad(loss, val) AS val FROM w2",
+        ])
+
+    On a [Context][xarray_sql.Context], each result is an
+    [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame].
+
+    Args:
+        ctx: Any DataFusion ``SessionContext``.
+        statements: SQL statements.
+
+    Returns:
+        One DataFrame per statement, in order.
+    """
+    if isinstance(statements, str):
+        raise TypeError("statements is a list of SQL statements")
+    frames: list = _ddx.sql_all(ctx, list(statements))
+    wrap = getattr(ctx, "_wrap", None)
+    return [wrap(f) for f in frames] if wrap else frames
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index b5f3a9c8..31dd1e85 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -25,6 +25,14 @@
 except ImportError:  # pragma: no cover - depends on the environment
     _proj = None
 
+_ad: ModuleType | None
+try:  # ddxdb is an optional dependency (`pip install xarray-sql[ml]`).
+    from . import ad
+
+    _ad = ad
+except ImportError:  # pragma: no cover - depends on the environment
+    _ad = None
+
 
 class Context(SessionContext):
     """A datafusion `SessionContext` that also supports `xarray.Dataset`s.
@@ -45,6 +53,10 @@ def __init__(self, *args, **kwargs):
         # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform.
         if _proj is not None:
             _proj.register(self)
+        # With ddxdb installed, .sql() also takes gradients: grad(), vjp()
+        # and jvp() in a select list, and ddx_stop_gradient().
+        if _ad is not None:
+            _ad.register(self)
 
     def from_dataset(
         self,
@@ -214,6 +226,10 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         ``.to_dataset(dimension_columns=[...])`` for round-tripping the
         result back to an ``xr.Dataset``.
 
+        With the ml extension installed (``pip install xarray-sql[ml]``),
+        the query may also call ``grad``, ``vjp`` and ``jvp`` (see
+        [xarray_sql.ad][]).
+
         Args:
             query: A SQL query string.
             *args: Forwarded to ``SessionContext.sql``.
@@ -222,7 +238,17 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         Returns:
             An [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapping the DataFusion DataFrame.
         """
-        inner = super().sql(query, *args, **kwargs)
+        if _ad is not None and _ad.calls_ad(query):
+            if args or kwargs:
+                raise TypeError(
+                    "a query that calls grad, vjp or jvp takes no other "
+                    "arguments"
+                )
+            return self._wrap(_ad.sql(self, query))
+        return self._wrap(super().sql(query, *args, **kwargs))
+
+    def _wrap(self, inner) -> XarrayDataFrame:
+        """Wrap a DataFusion DataFrame so it can round-trip to xarray."""
         return XarrayDataFrame(inner, templates=self._registered_datasets)