From 73157f2647554e736d33828536c609cc7f5187f2 Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:24:43 -0800
Subject: [PATCH 1/9] Reuse gradient programs and reduce Dataset registration
 overhead

Support prepared SQL gradients, arithmetic chunk-grid counts and direct Arrow registration. Preserve Context metadata through extension installation, expose target partitions and guard physical reuse against value-dependent filters.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 tests/test_ad.py      |  55 +++++++++++++++++--
 tests/test_prepare.py |  72 +++++++++++++++++++++++++
 xarray_sql/ad.py      |  50 +++++++++++++++++-
 xarray_sql/reader.py  |  12 ++---
 xarray_sql/sql.py     | 120 +++++++++++++++++++++++++++++++++++-------
 5 files changed, 278 insertions(+), 31 deletions(-)

diff --git a/tests/test_ad.py b/tests/test_ad.py
index 0f868bc7..df3d2675 100644
--- a/tests/test_ad.py
+++ b/tests/test_ad.py
@@ -136,9 +136,54 @@ def test_a_mention_in_a_literal_is_not_a_call():
     assert rows == [{"s": "grad(x)", "jvp(v)": 1}]
 
 
-def test_prepare_refuses_a_gradient(problem):
-    ctx, *_ = problem
-    with pytest.raises(NotImplementedError, match="grad"):
-        ctx.prepare(
-            f"WITH {LOSS} SELECT feature, grad(loss, w) AS g FROM weights"
+def test_context_prepare_runs_gradient_steps_and_reads_writes(problem):
+    ctx, x, w, y = problem
+    step = ctx.prepare(
+        f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature"
+    )
+    old = step()
+    original = old["w"].to_numpy().copy()
+    replacement = xr.Dataset(
+        {"w": ("feature", w.copy() * 2)}, coords={"feature": np.arange(3)}
+    )
+    got = step(weights=replacement)
+    np.testing.assert_allclose(
+        got["w"], 2 * x.T @ (x @ replacement.w.values - y)
+    )
+    np.testing.assert_array_equal(old["w"], original)
+    with pytest.raises(ValueError, match="coordinates"):
+        step(weights=replacement.assign_coords(feature=[3, 4, 5]))
+
+
+def test_prepared_gradients_read_new_values_and_keep_old_results(problem):
+    ctx, x, w, y = problem
+    statements = [
+        f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature",
+        f"WITH {LOSS} SELECT feature, w - 0.01 * grad(loss, w) AS w FROM weights ORDER BY feature",
+    ]
+    run = ad.prepare(ctx, statements)
+    [first, _] = run()
+    original = first["w"].to_numpy().copy()
+    for scale in (0.5, -1.0, 2.0):
+        values = xr.Dataset(
+            {"w": ("feature", w.copy() * scale)},
+            coords={"feature": np.arange(3)},
         )
+        gradient, step = run(weights=values)
+        want = 2 * x.T @ (x @ values.w.values - y)
+        np.testing.assert_allclose(gradient["w"], want)
+        np.testing.assert_allclose(step["w"], values.w.values - 0.01 * want)
+        np.testing.assert_array_equal(first["w"], original)
+    # Regular ddx preparation resolves a replacement table by name.
+    ctx.deregister_table("weights")
+    ctx.from_dataset("weights", values * 2)
+    [gradient, _] = run()
+    np.testing.assert_allclose(
+        gradient["w"], 2 * x.T @ (x @ (values.w.values * 2) - y)
+    )
+
+
+def test_prepared_statements_require_a_sequence(problem):
+    ctx, *_ = problem
+    with pytest.raises(TypeError, match="list"):
+        ad.prepare(ctx, "SELECT 1")
diff --git a/tests/test_prepare.py b/tests/test_prepare.py
index 8cc356bc..1b8b2b53 100644
--- a/tests/test_prepare.py
+++ b/tests/test_prepare.py
@@ -9,6 +9,11 @@
 
 from xarray_sql import Context
 
+try:
+    from einfold.datafusion import EinFold
+except ImportError:
+    EinFold = None
+
 # The build side of the join, `a`, is what a reused hash join would keep.
 JOIN = """
 SELECT b.i, SUM(t.v2 * b.w) AS s
@@ -17,6 +22,26 @@
 """
 
 
+@pytest.mark.parametrize("extension", [False, True])
+def test_join_key_data_columns_can_change_their_range(extension):
+    ctx = Context()
+    if extension:
+        if EinFold is None:
+            pytest.skip("requires the optional EinFold planner")
+        ctx = ctx.with_extensions(EinFold())
+    a = xr.Dataset({"key": ("i", np.arange(128)), "v": ("i", np.ones(128))})
+    b = xr.Dataset({"key": ("j", np.arange(256)), "w": ("j", np.arange(256.0))})
+    ctx.from_dataset("a", a)
+    ctx.from_dataset("b", b)
+    sql = "SELECT a.i,b.j,a.v+b.w AS s FROM a JOIN b ON a.key=b.key ORDER BY a.i,b.j"
+    query = ctx.prepare(sql)
+    for start in (0, 128, 16):
+        updated = a.copy(deep=True)
+        updated["key"][:] = np.arange(start, start + 128)
+        rows = query(a=updated)
+        assert rows.equals(ctx.sql(sql).to_arrow_table())
+
+
 def _a(v):
     return xr.Dataset(
         {"v": (("i", "j"), v)},
@@ -262,3 +287,50 @@ def test_a_quoted_table_replaced_in_sql_counts_too():
     ctx.sql('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)').collect()
     with pytest.raises(ValueError, match="registered again"):
         query()
+
+
+def test_extension_handle_keeps_dataset_metadata_and_invalidation(ctx):
+    want = ctx.sql("SELECT * FROM a").to_dataset()
+    derived = ctx.with_extensions()
+    query = derived.prepare("SELECT i, j, v FROM a ORDER BY i, j")
+    xr.testing.assert_identical(
+        derived.sql("SELECT * FROM a").to_dataset(), want
+    )
+    ctx.deregister_table("a")
+    with pytest.raises(ValueError, match="registered again"):
+        query()
+
+
+@pytest.mark.parametrize(
+    "chunks", [{"i": 1, "j": 1}, {"i": 1}, {"i": 4, "j": 5}]
+)
+def test_fine_chunk_registration_keeps_rows_metadata_and_writes(chunks):
+    values = np.arange(129 * 33, dtype=np.float64).reshape(129, 33)
+    ds = _a(values)
+    ds.attrs["experiment"] = "fine chunks"
+    ctx = Context()
+    ctx.from_dataset("a", ds, chunks=chunks)
+    sql = "SELECT i,j,v FROM a ORDER BY i,j"
+    xr.testing.assert_identical(ctx.sql(sql).to_dataset(), ds)
+    run = ctx.prepare(sql)
+    replacement = ds * 2
+    rows = run(a=replacement)
+    np.testing.assert_array_equal(rows["v"], replacement.v.values.reshape(-1))
+
+
+@pytest.mark.parametrize("threads", [1, 60])
+def test_thread_target_is_explicit_and_keeps_xarray_support(threads):
+    ctx = Context(threads=threads)
+    assert ctx.sql(
+        "SELECT value FROM information_schema.df_settings "
+        "WHERE name='datafusion.execution.target_partitions'"
+    ).to_pylist() == [{"value": str(threads)}]
+    ctx.from_dataset("a", _a(np.ones((4, 3))))
+    assert ctx.prepare("SELECT SUM(v) AS s FROM a")().to_pylist() == [
+        {"s": 12.0}
+    ]
+
+
+def test_thread_target_must_be_positive():
+    with pytest.raises(ValueError, match="positive"):
+        Context(threads=0)
diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py
index d0ff2c8f..eb12555a 100644
--- a/xarray_sql/ad.py
+++ b/xarray_sql/ad.py
@@ -43,7 +43,11 @@
 
 [sql_all][xarray_sql.ad.sql_all] runs several statements that take
 ``grad`` of the same loss for the price of one backward pass, which is
-how a training step updates each of a model's parameter tables. For the
+how a training step updates each of a model's parameter tables. For repeated
+training steps, [prepare][xarray_sql.ad.prepare] retains ddx's
+programs between calls and materializes each statement's result.
+[Context.prepare][xarray_sql.Context.prepare] handles one such statement;
+``ad.prepare`` handles several statements sharing a loss. For the
 programs underneath (``grad``, ``vjp``, ``jvp`` and ``run``, to build a
 backward pass once and run it every step), use ``ddxdb.ad`` directly;
 they take a ``Context`` like any other ``SessionContext``.
@@ -52,12 +56,13 @@
 from __future__ import annotations
 
 import re
+import threading
 from collections.abc import Sequence
 
 from datafusion import SessionContext
 from ddxdb import ad as _ddx
 
-__all__ = ["STOP_GRADIENT", "calls_ad", "register", "sql", "sql_all"]
+__all__ = ["STOP_GRADIENT", "calls_ad", "prepare", "register", "sql", "sql_all"]
 
 STOP_GRADIENT = _ddx.STOP_GRADIENT
 """The SQL name of the stop-gradient function, ``ddx_stop_gradient``."""
@@ -138,3 +143,44 @@ def sql_all(ctx: SessionContext, statements: Sequence[str]) -> list:
     frames: list = _ddx.sql_all(ctx, list(statements))
     wrap = getattr(ctx, "_wrap", None)
     return [wrap(f) for f in frames] if wrap else frames
+
+
+def prepare(ctx: SessionContext, statements: Sequence[str]):
+    """Prepare SQL gradient statements for repeated execution.
+
+    ddx retains the differentiation programs and checks table schemas,
+    views and session settings on each call. This avoids rebuilding those
+    programs in a training loop while allowing tables to be registered
+    again. Several statements can share a loss, as with ``sql_all``::
+
+        step = xql.ad.prepare(ctx, [
+            f"WITH loss AS ({LOSS}) "
+            "SELECT i, val - 0.01 * grad(loss, val) AS val FROM w ORDER BY i"
+        ])
+        [rows] = step()
+
+    On a ``Context``, a call may first write Dataset values into registered
+    tables, using the same validation as ``Context.prepare``::
+
+        [rows] = step(w=new_weights)
+
+    Returns a callable producing one materialized ``pyarrow.Table`` per
+    statement. Results remain valid after later calls. Writes and execution
+    are serialized for calls to this prepared function. Other users of the
+    context must coordinate concurrent writes themselves.
+    """
+    if isinstance(statements, str):
+        raise TypeError("statements is a list of SQL statements")
+    run = _ddx.prepare(ctx, list(statements))
+    lock = threading.Lock()
+
+    def execute(**tables):
+        with lock:
+            if tables:
+                write = getattr(ctx, "_write", None)
+                if write is None:
+                    raise TypeError("Dataset writes require xarray_sql.Context")
+                write(tables)
+            return [frame.to_arrow_table() for frame in run()]
+
+    return execute
diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py
index 8ec88bcf..db7c82c9 100644
--- a/xarray_sql/reader.py
+++ b/xarray_sql/reader.py
@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import itertools
+import math
 import os
 import sys
 from collections.abc import Callable, Iterable, Iterator
@@ -21,9 +22,9 @@
 import xarray as xr
 
 from .df import (
+    DEFAULT_BATCH_SIZE,
     Block,
     Chunks,
-    DEFAULT_BATCH_SIZE,
     _block_len,
     _block_metadata,
     _block_slices_from_resolved,
@@ -296,11 +297,10 @@ def _in_memory_partitions(
     # isn't contiguous in memory, and flattening it would copy. More
     # partitions than cores only adds overhead, since an in-memory table has
     # no partitions to prune.
+    # Count the Cartesian chunk grid without constructing slices that this
+    # in-memory path never reads, even for very finely chunked inputs.
     count = min(
-        sum(
-            1
-            for _ in _block_slices_from_resolved(ds, resolve_chunks(ds, chunks))
-        ),
+        math.prod(len(axis) for axis in resolve_chunks(ds, chunks).values()),
         os.cpu_count() or 1,
     )
     # `per_partition` batches of `size` rows each make `count` partitions,
@@ -324,7 +324,7 @@ def read_xarray_table(
     _iteration_callback: (
         Callable[[Block, list[str] | None], None] | None
     ) = None,
-) -> "LazyArrowStreamTable":
+) -> LazyArrowStreamTable:
     """Create a lazy DataFusion table from an xarray Dataset.
 
     This is the simplest way to register xarray data with DataFusion.
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index f64166a4..bf801903 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -1,6 +1,8 @@
 import re
 import time
 import uuid
+from collections.abc import Callable
+from types import ModuleType
 
 import numpy as np
 import pyarrow as pa
@@ -9,7 +11,6 @@
 from datafusion.catalog import Schema
 from datafusion.expr import TableScan
 from datafusion.plan import ExecutionPlan, LogicalPlan
-from types import ModuleType
 
 from . import cftime as cft
 from .df import (
@@ -38,14 +39,30 @@
 except ImportError:  # pragma: no cover - depends on the environment
     _ad = None
 
+_DYNAMIC_FILTERS = tuple(
+    f"datafusion.optimizer.enable{kind}_dynamic_filter_pushdown"
+    for kind in ("", "_join", "_topk", "_aggregate")
+)
+
 
 class Context(SessionContext):
     """A datafusion `SessionContext` that also supports `xarray.Dataset`s.
 
     Also available under its original name, ``XarrayContext``.
+    ``threads=`` sets DataFusion's target parallel partitions, overriding
+    that setting in an explicit ``config``. The default is the CPU count.
+    It is a concurrency target, rather than a fixed number of OS threads.
     """
 
-    def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
+    def __init__(
+        self,
+        config: SessionConfig | None = None,
+        *args,
+        threads: int | None = None,
+        **kwargs,
+    ):
+        if threads is not None and threads < 1:
+            raise ValueError("threads must be positive")
         if config is None:
             # Split a query's work across partitions only where enough rows
             # reach it, by DataFusion's estimate. DataFusion otherwise splits
@@ -65,6 +82,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
                     "true",
                 )
             )
+        if threads is not None:
+            config = config.with_target_partitions(threads)
         super().__init__(config, *args, **kwargs)
         # Track registered xarray Datasets so XarrayDataFrame can recover
         # defaults (dimension_columns) and metadata (var/dataset attrs,
@@ -78,6 +97,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
         # registered again from Python was deregistered first; SQL that
         # replaces or drops a table counts too (see sql()).
         self._deregistrations: dict[str, int] = {}
+        self._custom_udfs: set[str] = set()
+        self._trusted_math_udfs: set[str] = set()
         # With pyproj installed, every context speaks CRS out of the box:
         # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform.
         if _proj is not None:
@@ -87,6 +108,41 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
         if _ad is not None:
             _ad.register(self)
 
+    def register_udf(self, udf):
+        """Register a scalar function and track overrides for optional lowering."""
+        result = super().register_udf(udf)
+        self._custom_udfs.add(udf.name.lower())
+        self._trusted_math_udfs.discard(udf.name.lower())
+        return result
+
+    def register_udaf(self, udaf):
+        """Track aggregate-function overrides for optional numerical lowering."""
+        result = super().register_udaf(udaf)
+        self._custom_udfs.add(udaf.name.lower())
+        self._trusted_math_udfs.discard(udaf.name.lower())
+        return result
+
+    def register_udwf(self, udwf):
+        """Track window-function overrides for optional numerical lowering."""
+        result = super().register_udwf(udwf)
+        self._custom_udfs.add(udwf.name.lower())
+        self._trusted_math_udfs.discard(udwf.name.lower())
+        return result
+
+    def with_extensions(self, *extensions) -> "Context":
+        """Install DataFusion extension bundles, retaining xarray support.
+
+        Returns another handle on the same session, with this Context's
+        Dataset metadata and prepared-query invalidation state shared.
+        Extension codecs belong to the returned handle, so use it for
+        subsequent queries, as with DataFusion's own method.
+        """
+        session = super().with_extensions(*extensions)
+        derived = object.__new__(type(self))
+        derived.__dict__.update(self.__dict__)
+        derived.ctx = session.ctx
+        return derived
+
     def from_dataset(
         self,
         name: str,
@@ -214,9 +270,13 @@ def _from_dataset(
             table = read_xarray_table(
                 input_table, chunks, coord_arrays=coord_arrays
             )
+            register(table_name, table)
+        elif schema is None:
+            # Uniform Datasets can register their shared Arrow buffers
+            # directly, without a temporary catalog entry and lookup.
+            self.register_record_batches(table_name, partitions)
         else:
-            table = self._in_memory_table(partitions)
-        register(table_name, table)
+            register(table_name, self._in_memory_table(partitions))
         self._maybe_register_cftime_udf(input_table)
         return self
 
@@ -324,16 +384,11 @@ def prepare(self, query: str) -> "PreparedQuery":
             A [PreparedQuery][xarray_sql.PreparedQuery]; call it to run
             the query.
 
-        Raises:
-            NotImplementedError: If the query calls ``grad``, ``vjp`` or
-                ``jvp``, which only ``sql()`` computes for now.
+        Queries calling ``grad``, ``vjp`` or ``jvp`` retain ddx's programs
+        through [xarray_sql.ad.prepare][]. Their first call prepares the
+        programs; later calls check the catalog and settings and refresh
+        them if necessary.
         """
-        if _ad is not None and _ad.calls_ad(query):
-            # TODO(#294): prepare the query's programs with ddx.
-            raise NotImplementedError(
-                "prepare() doesn't compute grad, vjp or jvp yet; run the "
-                "query with sql(), or several with xarray_sql.ad.sql_all()"
-            )
         return PreparedQuery(self, query)
 
     def _write(self, writes: dict[str, xr.Dataset]) -> None:
@@ -415,10 +470,19 @@ class PreparedQuery:
     query again. A call raises if a table the query reads has been
     registered again since it was prepared. Functions such as ``now()``
     are evaluated once, when the query is prepared.
+
+    A query calling ``grad``, ``vjp`` or ``jvp`` instead retains ddx's
+    programs. Its first call prepares them, and subsequent calls refresh
+    them when schemas, views or settings change. Tables may be registered
+    again under the same names, as with [xarray_sql.ad.prepare][].
     """
 
     def __init__(self, ctx: Context, query: str):
         self._ctx = ctx
+        self._gradient: Callable[..., list[pa.Table]] | None = None
+        if _ad is not None and _ad.calls_ad(query):
+            self._gradient = _ad.prepare(ctx, [query])
+            return
         start = time.perf_counter()
         df = SessionContext.sql(ctx, query)
         plan = df.execution_plan()
@@ -441,11 +505,13 @@ def __init__(self, ctx: Context, query: str):
         }
         # A plan keeps state from a run (a hash join keeps the table it
         # built), so each call runs a fresh copy of it, decoded from its
-        # serialized form. Copying a plan copies its in-memory tables, so
-        # over large ones it's slower than planning again; a plan that reads
-        # a lazy table, or is nested deeper than protobuf decodes, can't be
-        # copied at all. Either way, each call plans the query again.
+        # serialized form. The default codec copies in-memory tables, so
+        # over large ones it can be slower than planning again; extensions
+        # can supply a codec that resets state while retaining input buffers.
+        # A plan the session's codecs cannot copy is planned again each call.
         self._plan: ExecutionPlan | None = None
+        if not _safe_physical_reuse(ctx):
+            return
         start = time.perf_counter()
         try:
             _copy(ctx, plan)
@@ -464,6 +530,9 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table:
         Returns:
             The query's rows.
         """
+        if self._gradient is not None:
+            (rows,) = self._gradient(**tables)
+            return rows
         for name in tables:
             if name not in self._registered:
                 raise ValueError(
@@ -495,7 +564,22 @@ def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan:
     # TODO(#283): reset the plan's state with datafusion-python's binding of
     # reset_plan_states, once it has one, rather than decode it again.
     # `plan` itself never runs, so it has no state to leave behind.
-    return ExecutionPlan.from_bytes(ctx, plan.to_bytes())
+    return ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx))
+
+
+def _safe_physical_reuse(ctx: SessionContext) -> bool:
+    """Dynamic filters retain value-dependent bounds even after a reset."""
+    names = ", ".join(repr(name) for name in _DYNAMIC_FILTERS)
+    try:
+        rows = SessionContext.sql(
+            ctx,
+            f"SELECT value FROM information_schema.df_settings WHERE name IN ({names})",
+        ).to_arrow_table()
+    except ValueError:
+        return False
+    return rows.num_rows == len(_DYNAMIC_FILTERS) and all(
+        value == "false" for value in rows["value"].to_pylist()
+    )
 
 
 def _tables_read(plan: LogicalPlan) -> set[str]:

From 9c847894c79484c281d10b2aa1277d054a0d5f05 Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:31:04 -0800
Subject: [PATCH 2/9] Share SQL stages and dispatch validated numerical kernels

Materialize shared reachable CTEs once, reuse stage programs with fresh physical state, and retain ddx differentiation of original views. Add optional session-local float64 CPU and float32 MLX kernel execution with SQL validation and fallback, using the separately maintained EinFold feature build.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 Cargo.lock                        |   1 +
 Cargo.toml                        |   1 +
 accelerators/einfold_mlx/build.py |  97 ++++
 benchmarks/ml/01_attention.py     |  24 +-
 benchmarks/ml/02_nanogpt.py       | 284 ++++++++++--
 benchmarks/ml/_bench.py           |  20 +-
 benchmarks/ml/_hosts.py           | 309 ++++++++++++-
 pyproject.toml                    |   8 +-
 src/lib.rs                        |   3 +
 src/shared.rs                     | 137 ++++++
 tests/test_ml_demos.py            |   7 +-
 tests/test_mlx.py                 | 391 +++++++++++++++++
 tests/test_shared.py              | 296 +++++++++++++
 uv.lock                           |  69 ++-
 xarray_sql/mlx.py                 | 258 +++++++++++
 xarray_sql/shared.py              | 708 ++++++++++++++++++++++++++++++
 16 files changed, 2541 insertions(+), 72 deletions(-)
 create mode 100644 accelerators/einfold_mlx/build.py
 create mode 100644 src/shared.rs
 create mode 100644 tests/test_mlx.py
 create mode 100644 tests/test_shared.py
 create mode 100644 xarray_sql/mlx.py
 create mode 100644 xarray_sql/shared.py

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

From 496cb11e4d584f4433f2ee2db53cc74d35ff273d Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:34:37 -0800
Subject: [PATCH 3/9] Pin the separately published EinFold kernel bridge

Build the upstream-compatible feature commit from the contribution fork.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 accelerators/einfold_mlx/build.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/accelerators/einfold_mlx/build.py b/accelerators/einfold_mlx/build.py
index 10e869f0..7991034f 100644
--- a/accelerators/einfold_mlx/build.py
+++ b/accelerators/einfold_mlx/build.py
@@ -17,7 +17,7 @@
 import tempfile
 from pathlib import Path
 
-_REVISION = "09a75e24c07efb320a5c2c288d7c958021ae3a39"
+_REVISION = "0b2f6db892c3cbf9df2c3d775fbcf4c5654df093"
 
 
 def main():

From da24b305211513c71870acb139d82d24a079f4f3 Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:31:04 -0800
Subject: [PATCH 4/9] Retain compact arrays across numerical SQL programs

Lower supported DataFusion and ddx stages through EinFold compact programs, importing live inputs on each ordinary query and expanding coordinates only at Arrow snapshot boundaries. Bound computed intermediates, preserve row/NULL masks and retain conservative SQL fallback.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 tests/test_dense.py      | 283 ++++++++++++++++++++++++++++++
 xarray_sql/_dense_mlx.py | 365 +++++++++++++++++++++++++++++++++++++++
 xarray_sql/dense.py      | 277 +++++++++++++++++++++++++++++
 3 files changed, 925 insertions(+)
 create mode 100644 tests/test_dense.py
 create mode 100644 xarray_sql/_dense_mlx.py
 create mode 100644 xarray_sql/dense.py

diff --git a/tests/test_dense.py b/tests/test_dense.py
new file mode 100644
index 00000000..43287ab3
--- /dev/null
+++ b/tests/test_dense.py
@@ -0,0 +1,283 @@
+"""Compact execution preserves SQL rows and live Dataset inputs."""
+
+import os
+
+import numpy as np
+import pyarrow as pa
+import pytest
+import xarray as xr
+from datafusion import udf
+
+from xarray_sql import Context
+
+pytest.importorskip("ddxdb")
+pytest.importorskip("einfold.dense_program")
+from ddxdb import UnsupportedExpression
+from einfold.datafusion import EinFold
+
+from xarray_sql.dense import DenseExecutor
+
+_PRECISIONS = [
+    "float64",
+    pytest.param(
+        "float32",
+        marks=pytest.mark.skipif(
+            os.environ.get("MLX_ENABLE_TF32") != "0",
+            reason="GPU tests require MLX_ENABLE_TF32=0",
+        ),
+    ),
+]
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_conditional_contraction_stays_within_array_bound(precision):
+    ctx = Context()
+    left = np.arange(64, dtype=float).reshape(8, 8) / 64
+    right = np.arange(64, dtype=float).reshape(8, 8) / 32
+    ctx.from_dataset("a", xr.Dataset({"v": (("i", "k"), left)}))
+    ctx.from_dataset("b", xr.Dataset({"v": (("k", "j"), right)}))
+    executor = DenseExecutor(
+        ctx, precision=precision, fallback=False, max_array_bytes=2048
+    )
+    rows = executor.prepare("""WITH products AS (
+      SELECT a.i,b.j,a.k,a.v*b.v AS p FROM a JOIN b ON a.k=b.k)
+      SELECT i,j,SUM(CASE WHEN i>=j THEN p ELSE 0.0 END) AS x,
+                 SUM(CASE WHEN i>=j THEN 0.0 ELSE p END) AS y
+      FROM products GROUP BY i,j ORDER BY i,j""")()
+    result = left @ right
+    i, j = np.indices(result.shape)
+    np.testing.assert_allclose(
+        np.asarray(rows["x"]).reshape(8, 8), np.where(i >= j, result, 0)
+    )
+    np.testing.assert_allclose(
+        np.asarray(rows["y"]).reshape(8, 8), np.where(i >= j, 0, result)
+    )
+
+
+@pytest.fixture(params=_PRECISIONS)
+def executor(request):
+    ctx = Context().with_extensions(EinFold())
+    ctx.from_dataset(
+        "a",
+        xr.Dataset(
+            {"v": (("i", "j"), np.arange(6, dtype=float).reshape(2, 3))},
+            coords={"i": [-2, 7], "j": [8, 2, 5]},
+        ),
+    )
+    ctx.from_dataset(
+        "b",
+        xr.Dataset(
+            {"w": (("j", "k"), np.arange(6, dtype=float).reshape(3, 2))},
+            coords={"j": [5, 8, 2], "k": [9, 1]},
+        ),
+    )
+    return DenseExecutor(ctx, precision=request.param, fallback=False)
+
+
+@pytest.mark.parametrize(
+    "sql",
+    [
+        "SELECT a.i,b.k,SUM(a.v*b.w) AS v FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k ORDER BY a.i,b.k",
+        "SELECT i,AVG(v) AS m,COUNT(v) AS n FROM a WHERE v>1 GROUP BY i ORDER BY i",
+        "SELECT SUM(v) AS s,AVG(v) AS m,COUNT(v) AS n FROM a WHERE v<0",
+        "SELECT i,SUM(v) AS s FROM a WHERE v<0 GROUP BY i ORDER BY i",
+        "SELECT i,j,MAX(v) OVER(PARTITION BY i) AS m FROM a ORDER BY i,j",
+        "SELECT i,j,i/3 AS q,i%3 AS r FROM a ORDER BY i,j",
+        "SELECT i,j,(v-3.5)%2.0 AS r FROM a ORDER BY i,j",
+        "SELECT i,j,power(2,63) AS a,power(2,-1) AS b FROM a ORDER BY i,j",
+        "SELECT i,j,(CASE WHEN v>2 THEN TRUE END) AND FALSE AS a,(CASE WHEN v>2 THEN TRUE END) OR TRUE AS b FROM a ORDER BY i,j",
+        "WITH t AS (SELECT i,CASE WHEN v>2 THEN v END AS v FROM a) SELECT i,SUM(v) AS s,COUNT(v) AS n FROM t GROUP BY i ORDER BY i",
+        "WITH t AS (SELECT i,v FROM a UNION ALL SELECT i,2*v AS v FROM a) SELECT i,SUM(v) AS s,COUNT(v) AS n FROM t GROUP BY i ORDER BY i",
+        "SELECT a.i,a.j,b.w FROM a LEFT JOIN (SELECT j,k,w FROM b WHERE k=1) b ON a.j=b.j ORDER BY a.i,a.j",
+        "SELECT a.i,a.j,c.k FROM a LEFT JOIN (SELECT j,k,w FROM b WHERE k=1 AND w>1) b ON a.j=b.j JOIN (SELECT j,k FROM b WHERE k=9) c ON b.j=c.j ORDER BY a.i,a.j,c.k",
+        "SELECT a.i,a.j,c.k FROM a LEFT JOIN (SELECT j,k,w FROM b WHERE k=1 AND w>1) b ON a.j=b.j JOIN (SELECT j,k FROM b WHERE k=9) c ON b.j IS NOT DISTINCT FROM c.j ORDER BY a.i,a.j,c.k",
+    ],
+)
+def test_dense_sql_rows_match_datafusion(executor, sql):
+    actual = executor.prepare(sql)()
+    expected = executor.ctx.sql(sql).to_arrow_table()
+    assert actual.column_names == expected.column_names
+    for a, e in zip(actual.columns, expected.columns):
+        assert a.null_count == e.null_count
+        assert a.to_pylist() == e.to_pylist()
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_dense_gather_gradient_and_snapshots(precision):
+    ctx = Context().with_extensions(EinFold())
+    weights = xr.Dataset(
+        {"v": (("tok", "e"), np.arange(6, dtype=float).reshape(3, 2))},
+        coords={"tok": [8, -1, 4]},
+    )
+    ctx.from_dataset("w", weights)
+    ctx.from_dataset("batch", xr.Dataset({"tok": ("pos", [4, 4, 8])}))
+    executor = DenseExecutor(ctx, precision=precision, fallback=False)
+    executor.register_view(
+        "model", "SELECT b.pos,w.e,w.v FROM batch b JOIN w ON b.tok=w.tok"
+    )
+    run = executor.prepare(
+        "WITH loss AS (SELECT SUM(v*v)/2 AS l FROM model) SELECT tok,e,grad(loss,v) AS g FROM w ORDER BY tok,e"
+    )
+    first = run()
+    assert first.to_pydict() == {
+        "tok": [-1, -1, 4, 4, 8, 8],
+        "e": [0, 1, 0, 1, 0, 1],
+        "g": [0.0, 0.0, 8.0, 10.0, 0.0, 1.0],
+    }
+    weights.v.values[:] += 2
+    second = run()
+    assert second["g"].to_pylist() == [0.0, 0.0, 12.0, 14.0, 2.0, 3.0]
+    assert first["g"].to_pylist() == [0.0, 0.0, 8.0, 10.0, 0.0, 1.0]
+
+
+def test_dense_falls_back_for_overridden_scalar_function(executor):
+    ctx = executor.ctx
+    ctx.register_udf(
+        udf(
+            lambda a: pa.array(np.asarray(a) + 100),
+            [pa.float64()],
+            pa.float64(),
+            "immutable",
+            "exp",
+        )
+    )
+    run = DenseExecutor(ctx).prepare(
+        "SELECT i,j,exp(v) AS v FROM a ORDER BY i,j"
+    )
+    assert run().equals(
+        ctx.sql("SELECT i,j,exp(v) AS v FROM a ORDER BY i,j").to_arrow_table()
+    )
+    ctx.register_udf(
+        udf(
+            lambda a: pa.array(np.asarray(a) + 200),
+            [pa.float64()],
+            pa.float64(),
+            "immutable",
+            "exp",
+        )
+    )
+    assert run().equals(
+        ctx.sql("SELECT i,j,exp(v) AS v FROM a ORDER BY i,j").to_arrow_table()
+    )
+
+
+def test_dense_coordinate_cast_retains_its_integer_width():
+    ctx = Context().with_extensions(EinFold())
+    ctx.from_dataset(
+        "a",
+        xr.Dataset(
+            {"v": ("i", [1.0, 2.0])},
+            coords={"i": np.array([2000000000, 2000000001], dtype=np.int32)},
+        ),
+    )
+    sql = "SELECT i,CAST(i AS BIGINT)*2 AS j FROM a ORDER BY i"
+    assert (
+        DenseExecutor(ctx, fallback=False)
+        .prepare(sql)()
+        .equals(ctx.sql(sql).to_arrow_table())
+    )
+
+
+def test_dense_base_log_fallback_never_writes_an_input():
+    ctx = Context().with_extensions(EinFold())
+    values = np.array([1.0, 2.0, 3.0])
+    ctx.from_dataset("x", xr.Dataset({"v": ("i", values)}))
+    sql = "SELECT i,log(2.,v) AS v FROM x ORDER BY i"
+    expected = ctx.sql(sql).to_arrow_table()
+    assert DenseExecutor(ctx).prepare(sql)().equals(expected)
+    np.testing.assert_array_equal(values, [1.0, 2.0, 3.0])
+
+
+def test_dense_integer_sum_uses_the_sql_accumulator_width():
+    ctx = Context().with_extensions(EinFold())
+    ctx.from_dataset(
+        "x",
+        xr.Dataset(
+            {"v": ("i", np.array([2000000000, 2000000001], dtype=np.int32))}
+        ),
+    )
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_dense_null_and_empty_filters_do_not_evaluate_deleted_rows(precision):
+    ctx = Context().with_extensions(EinFold())
+    ctx.from_dataset("x", xr.Dataset({"v": ("i", [1.0, 2.0, 3.0])}))
+    executor = DenseExecutor(ctx, precision=precision, fallback=False)
+    for sql in [
+        "SELECT i,v FROM x WHERE i=NULL ORDER BY i",
+        "SELECT i,1/(i-i) AS v FROM x WHERE i<0 ORDER BY i",
+    ]:
+        assert executor.prepare(sql)().equals(ctx.sql(sql).to_arrow_table())
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_dense_integer_averages_with_repeated_group_keys(precision):
+    ctx = Context().with_extensions(EinFold())
+    ctx.from_dataset(
+        "x", xr.Dataset({"key": ("i", [1, 1, 2, 2]), "v": ("i", [1, 2, 3, 4])})
+    )
+    sql = "SELECT key,AVG(v) AS m,SUM(v) AS s FROM x GROUP BY key ORDER BY key"
+    assert (
+        DenseExecutor(ctx, precision=precision, fallback=False)
+        .prepare(sql)()
+        .equals(ctx.sql(sql).to_arrow_table())
+    )
+    sql = "SELECT SUM(v) AS s,AVG(v) AS m FROM x"
+    assert (
+        DenseExecutor(ctx, fallback=False)
+        .prepare(sql)()
+        .equals(ctx.sql(sql).to_arrow_table())
+    )
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_dense_unsupported_backward_region_restarts_with_original_sql(
+    precision,
+):
+    ctx = Context().with_extensions(EinFold())
+    values = np.arange(6.0, dtype=float).reshape(2, 3)
+    ctx.from_dataset("a", xr.Dataset({"v": (("i", "j"), values)}))
+    query = DenseExecutor(ctx, precision=precision).prepare(
+        "WITH t AS (SELECT CAST(i AS DOUBLE) AS i,CAST(j AS DOUBLE) AS j,SUM(v*v) AS s FROM a GROUP BY CAST(i AS DOUBLE),CAST(j AS DOUBLE)),loss AS (SELECT SUM(s) AS l FROM t) SELECT i,j,grad(loss,v) AS g FROM a ORDER BY i,j"
+    )
+    assert query()["g"].to_pylist() == [0.0, 2.0, 4.0, 6.0, 8.0, 10.0]
+    values[:] += 1
+    assert query()["g"].to_pylist() == [2.0, 4.0, 6.0, 8.0, 10.0, 12.0]
+
+
+def test_dense_directional_derivatives_and_typed_refusal(executor):
+    ctx = executor.ctx
+    tangent = "WITH tangent AS (SELECT j,k,CAST(k AS DOUBLE) AS w FROM b) "
+    rows = executor.prepare(
+        tangent
+        + "SELECT a.i,b.k,jvp(SUM(a.v*b.w),b.w,tangent) AS p FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k ORDER BY i,k"
+    )()
+    assert rows.to_pydict() == {
+        "i": [-2, -2, 7, 7],
+        "k": [1, 9, 1, 9],
+        "p": [3.0, 27.0, 12.0, 108.0],
+    }
+    ctx.from_dataset(
+        "cotangent",
+        xr.Dataset(
+            {"v": (("i", "k"), [[1.0, 2.0], [3.0, 4.0]])},
+            coords={"i": [-2, 7], "k": [9, 1]},
+        ),
+    )
+    predicted = "WITH pred AS (SELECT a.i,b.k,SUM(a.v*b.w) AS v FROM a JOIN b ON a.j=b.j GROUP BY a.i,b.k) "
+    rows = executor.prepare(
+        predicted + "SELECT j,k,vjp(pred,w,cotangent) AS g FROM b ORDER BY j,k"
+    )()
+    assert rows.to_pydict() == {
+        "j": [2, 2, 5, 5, 8, 8],
+        "k": [1, 9, 1, 9, 1, 9],
+        "g": [18.0, 13.0, 24.0, 17.0, 12.0, 9.0],
+    }
+    rows = executor.prepare(
+        "WITH loss AS (SELECT SUM(v*ddx_stop_gradient(v)) AS l FROM a) SELECT i,j,grad(loss,v) AS g FROM a ORDER BY i,j"
+    )()
+    assert rows["g"].to_pylist() == [1.0, 2.0, 0.0, 4.0, 5.0, 3.0]
+    with pytest.raises(UnsupportedExpression, match="GROUP BY key"):
+        executor.prepare(
+            "WITH buckets AS (SELECT ROUND(v) AS k,SUM(v*v) AS s FROM a GROUP BY ROUND(v)),loss AS (SELECT SUM(k*s) AS l FROM buckets) SELECT i,j,grad(loss,v) AS g FROM a"
+        )()
diff --git a/xarray_sql/_dense_mlx.py b/xarray_sql/_dense_mlx.py
new file mode 100644
index 00000000..1ffe642f
--- /dev/null
+++ b/xarray_sql/_dense_mlx.py
@@ -0,0 +1,365 @@
+"""MLX array operations for EinFold's compact numerical programs.
+
+[MLX](https://github.com/ml-explore/mlx) evaluates the arithmetic chosen by
+[EinFold](https://github.com/xqlsystems/einfold). Compiled functions take
+every numerical leaf as an argument, so a cached function never captures
+parameter values from an earlier training step. Coordinates stay on the CPU.
+"""
+
+from __future__ import annotations
+
+import os
+import string
+from collections import Counter
+
+import mlx.core as mx
+import numpy as np
+from einfold.dense_program import _OPS, _align, _Decline, _op
+
+
+def _remainder(a, b):
+    result = mx.remainder(a, b)
+    return mx.where((result != 0) & ((a < 0) != (b < 0)), result - b, result)
+
+
+_FUNCTIONS = {
+    "+": mx.add,
+    "-": mx.subtract,
+    "*": mx.multiply,
+    "/": mx.divide,
+    "%": _remainder,
+    "=": mx.equal,
+    "!=": mx.not_equal,
+    "<": mx.less,
+    "<=": mx.less_equal,
+    ">": mx.greater,
+    ">=": mx.greater_equal,
+    "and": mx.logical_and,
+    "or": mx.logical_or,
+    "not": mx.logical_not,
+    "neg": mx.negative,
+    "where": mx.where,
+    "abs": mx.abs,
+    "sqrt": mx.sqrt,
+    "exp": mx.exp,
+    "erf": mx.erf,
+    "ln": mx.log,
+    "log": mx.log10,
+    "tanh": mx.tanh,
+    "isnan": mx.isnan,
+    "iszero": lambda x: x == 0,
+    "signum": mx.sign,
+    "sign": mx.sign,
+    "sin": mx.sin,
+    "cos": mx.cos,
+    "power": mx.power,
+    "pow": mx.power,
+    "floor": mx.floor,
+    "ceil": mx.ceil,
+    "minimum": mx.minimum,
+    "maximum": mx.maximum,
+    "float64": lambda x: x.astype(mx.float32),
+    "int64": lambda x: x.astype(mx.int64),
+    "bool": lambda x: x.astype(mx.bool_),
+}
+
+
+class _Runtime:
+    def __init__(self):
+        if os.environ.get("MLX_ENABLE_TF32") != "0":
+            raise ValueError("launch Python with MLX_ENABLE_TF32=0")
+        self.compiled = {}
+        self.counts = Counter()
+        self.tracing = False
+
+    def owns(self, array):
+        return isinstance(array, mx.array)
+
+    def boolean(self, value):
+        self.counts["gpu_status_readbacks"] += 1
+        return bool(value.item())
+
+    def status(self, values, flags):
+        status = mx.array(True)
+        for flag in flags:
+            status = mx.logical_and(status, mx.all(self._mx(flag)))
+        for value in values:
+            status = mx.logical_and(
+                status, mx.all(mx.isfinite(self._mx(value)))
+            )
+        return status
+
+    def validate(self, values, flags):
+        status = self.status(values, flags)
+        mx.eval(values, status)
+        if not self.boolean(status):
+            raise ValueError(
+                "training update must be rectangular, nonnullable and finite"
+            )
+
+    def compile_training(self, function):
+        self.counts["gpu_training_compilations"] += 1
+        return mx.compile(function)
+
+    def finish_training(self, arrays, status):
+        self.counts["gpu_training_executions"] += 1
+        mx.eval(arrays, status)
+        if not self.boolean(status):
+            raise ValueError(
+                "training update must be rectangular, nonnullable and finite"
+            )
+
+    def indices(self, values):
+        if values.size and (
+            int(values.min()) < -(1 << 63) or int(values.max()) >= (1 << 63)
+        ):
+            raise _Decline("device index is outside exact int64 range")
+        self.counts["gpu_input_arrays"] += 1
+        self.counts["gpu_input_bytes"] += values.nbytes
+        return mx.array(values, dtype=mx.int64)
+
+    def random_key(self, seed, step):
+        return mx.bitwise_xor(
+            mx.random.key(seed), mx.array([0, step], dtype=mx.uint32)
+        )
+
+    def dropout(self, key, shapes, probability):
+        keys = mx.random.split(key, len(shapes))
+        return [
+            mx.random.bernoulli(1 - probability, shape=shape, key=k).astype(
+                mx.float32
+            )
+            / (1 - probability)
+            for k, shape in zip(keys, shapes)
+        ]
+
+    def memory(self):
+        return {
+            "active_bytes": mx.get_active_memory(),
+            "cache_bytes": mx.get_cache_memory(),
+            "peak_bytes": mx.get_peak_memory(),
+        }
+
+    def copy(self, value):
+        return mx.array(value)
+
+    def numpy(self, value):
+        if self.owns(value):
+            if self.tracing:
+                raise _Decline(
+                    "compiled numerical graph requires a host conversion"
+                )
+            self.counts["gpu_readbacks"] += 1
+            self.counts["gpu_readback_bytes"] += value.size * value.itemsize
+        return np.asarray(value)
+
+    def array(self, value):
+        array = np.asarray(value)
+        if array.dtype.kind == "f":
+            # ddx uses the extreme float64 value as a masked MAX sentinel.
+            # It stays beyond every supported finite float32 input. An
+            # observable infinity is rejected at the final SQL boundary.
+            if array.ndim == 0 and abs(array) == np.finfo(np.float64).max:
+                return mx.array(float("-inf") if array < 0 else float("inf"))
+            finite = np.isfinite(array)
+            if (array.ndim and not finite.all()) or np.any(
+                np.abs(array[finite]) > np.finfo(np.float32).max
+            ):
+                raise _Decline(
+                    f"GPU input outside finite float32 range: {array.dtype}, shape={array.shape}, finite={finite.sum()}/{array.size}"
+                )
+            self.counts["gpu_input_arrays"] += 1
+            self.counts["gpu_input_bytes"] += array.nbytes
+            return mx.array(array, dtype=mx.float32)
+        return array
+
+    def _mx(self, value):
+        return value if self.owns(value) else mx.array(value)
+
+    def pointwise(self, tensor):
+        leaves, instructions, seen = [], [], {}
+
+        def describe(node):
+            if node in seen:
+                return seen[node]
+            if node.kind in _OPS and node.cached is None and len(seen) < 64:
+                children = tuple(describe(a) for a in node.args)
+                instruction = node.kind, children
+            else:
+                leaves.append(_align(node.get(), node.axes, tensor.axes))
+                instruction = "leaf", len(leaves) - 1
+            at = len(instructions)
+            seen[node] = at
+            instructions.append(instruction)
+            return at
+
+        describe(tensor)
+        signature = tuple(instructions)
+        if not any(self.owns(a) for a in leaves):
+            return None
+        if signature not in self.compiled:
+
+            def execute(*args):
+                values = []
+                for kind, children in signature:
+                    if kind == "leaf":
+                        values.append(args[children])
+                    else:
+                        values.append(
+                            _FUNCTIONS[kind](*(values[c] for c in children))
+                        )
+                return values[-1]
+
+            self.compiled[signature] = mx.compile(execute)
+            self.counts["gpu_fused_programs"] += 1
+        self.counts["gpu_fused_executions"] += 1
+        return self.compiled[signature](*(self._mx(a) for a in leaves))
+
+    def take(self, data, indexes, axis):
+        return (
+            mx.take(self._mx(data), self._mx(indexes), axis=axis)
+            if self.owns(data) or self.owns(indexes)
+            else np.take(data, indexes, axis=axis)
+        )
+
+    def indexes(self, domain, values):
+        if domain.dtype.kind not in "iu" or not len(domain):
+            raise _Decline("GPU index domain must contain integers")
+        if int(domain.min()) < -(1 << 63) or int(domain.max()) >= (1 << 63):
+            raise _Decline("GPU index domain is outside exact int64 range")
+        domain = domain.astype(np.int64, copy=False)
+        order = np.argsort(domain)
+        sorted_domain = mx.array(domain[order])
+        at = mx.minimum(mx.searchsorted(sorted_domain, values), len(domain) - 1)
+        indexes = mx.take(mx.array(order), at)
+        return indexes, mx.take(mx.array(domain), indexes) == values
+
+    def index(self, data, indexes):
+        return (
+            self._mx(data)[tuple(self._mx(i) for i in indexes)]
+            if self.owns(data) or any(self.owns(i) for i in indexes)
+            else data[tuple(indexes)]
+        )
+
+    def contract(self, arrays, labels, output):
+        converted = [self._mx(a) for a in arrays]
+        if not any(mx.issubdtype(a.dtype, mx.floating) for a in converted):
+            for array in arrays:
+                if self.owns(array):
+                    if array.dtype != mx.bool_:
+                        raise _Decline(
+                            "GPU integer contraction requires a count mask"
+                        )
+                elif np.any((array < -1) | (array > 1)):
+                    raise _Decline(
+                        "GPU integer contraction requires unit factors"
+                    )
+            dimensions = {
+                label: size
+                for a, names in zip(converted, labels)
+                for label, size in zip(names, a.shape)
+            }
+            used = tuple(sorted(dimensions))
+            value = mx.array(1, dtype=mx.int64)
+            for a, names in zip(converted, labels):
+                ordered = [n for n in used if n in names]
+                a = a.transpose(tuple(names.index(n) for n in ordered))
+                a = a.reshape(
+                    tuple(dimensions[n] if n in names else 1 for n in used)
+                )
+                value = value * a.astype(mx.int64)
+            value = mx.sum(
+                value,
+                axis=tuple(i for i, n in enumerate(used) if n not in output),
+            )
+            remaining = [n for n in used if n in output]
+            return value.transpose(tuple(remaining.index(n) for n in output))
+        equation = (
+            ",".join(
+                "".join(string.ascii_letters[i] for i in label)
+                for label in labels
+            )
+            + "->"
+            + "".join(string.ascii_letters[i] for i in output)
+        )
+        # Mask/count operands are promoted before contraction; a boolean
+        # reduction otherwise counts presence rather than SQL rows.
+        arrays = [a.astype(mx.float32) for a in converted]
+        self.counts["gpu_contractions"] += 1
+        return mx.einsum(equation, *arrays)
+
+    def broadcast(self, data, source, target, shape):
+        return mx.broadcast_to(_align(self._mx(data), source, target), shape)
+
+    def extreme(self, kind, data, source, mask, mask_axes, used, kept):
+        shape = tuple(len(a) for a in used)
+        if any(size == 0 for size in shape):
+            raise _Decline("empty GPU extremum")
+        data = self.broadcast(data, source, used, shape)
+        mask = self.broadcast(mask, mask_axes, used, shape)
+        data = mx.where(
+            mask, data, -float("inf") if kind == "max" else float("inf")
+        )
+        axes = tuple(i for i, a in enumerate(used) if a not in kept)
+        return (mx.max if kind == "max" else mx.min)(data, axis=axes)
+
+    def scatter(self, data, indexes, size, *, count_bound):
+        data = self._mx(data)
+        integer = data.dtype == mx.int64
+        if integer:
+            if count_bound > np.iinfo(np.int32).max:
+                raise _Decline("GPU integer scatter exceeds exact int32 range")
+            data = data.astype(mx.int32)
+        result = mx.zeros((data.shape[0], size), dtype=data.dtype)
+        self.counts["gpu_scatter_reductions"] += 1
+        result = result.at[
+            mx.arange(data.shape[0])[:, None], mx.array(indexes)
+        ].add(data)
+        return result.astype(mx.int64) if integer else result
+
+    def merge(self, parts, keys, key, axes, *, flags=None):
+        result = mx.zeros(tuple(len(a) for a in axes), dtype=mx.float32)
+        for part in parts:
+            source = tuple(part.columns[k].tensor.axis for k in keys)
+            indexes = [
+                np.searchsorted(a.values, s.values)
+                for a, s in zip(axes, source)
+            ]
+            shape = tuple(len(a) for a in source)
+            if flags is not None:
+                mask = (
+                    part.present
+                    if flags == "rows"
+                    else _op("and", part.present, part.columns[key].valid)
+                )
+                data = mx.ones(shape, dtype=mx.float32)
+            else:
+                col = part.columns[key]
+                value = col.tensor.get()
+                if self.owns(value) and not mx.issubdtype(
+                    value.dtype, mx.floating
+                ):
+                    raise _Decline(
+                        "GPU integer partial aggregate requires relational execution"
+                    )
+                mask = _op("and", part.present, col.valid)
+                data = self.broadcast(value, col.tensor.axes, source, shape)
+            valid = self.broadcast(mask.get(), mask.axes, source, shape)
+            data = mx.where(valid, data, 0)
+            if axes:
+                grid = [
+                    mx.array(
+                        i.reshape(
+                            tuple(
+                                len(i) if j == at else 1
+                                for j in range(len(axes))
+                            )
+                        )
+                    )
+                    for at, i in enumerate(indexes)
+                ]
+                result = result.at[tuple(grid)].add(data)
+            else:
+                result = result + data
+        self.counts["gpu_union_reductions"] += 1
+        return result
diff --git a/xarray_sql/dense.py b/xarray_sql/dense.py
new file mode 100644
index 00000000..f3ef3685
--- /dev/null
+++ b/xarray_sql/dense.py
@@ -0,0 +1,277 @@
+"""Experimental compact execution of SQL numerical programs.
+
+[EinFold](https://github.com/xqlsystems/einfold) lowers DataFusion plans to
+array operations. Its optional ``dense_program`` module retains compact
+arrays between shared SQL and [ddx](https://github.com/xqlsystems/ddx)
+automatic differentiation stages. Public results remain Arrow tables.
+The original SQL executor handles unsupported plans without changing SQL
+or gradient definitions. Install the patch in ``accelerators/einfold_mlx``.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Any
+
+import pyarrow as pa
+from datafusion import SessionContext
+from ddxdb import ad as _ddx
+from ddxdb import prepare as _prepare_ad
+from einfold.dense_program import (
+    _CACHE,
+    _EXECUTION,
+    _MAX_BYTES,
+    _OPS,
+    DenseProgram,
+    _ComputedCache,
+    _Decline,
+    _read_dataset,
+)
+
+from xarray_sql.ad import calls_ad
+from xarray_sql.shared import SharedExecutor, _Backend, _Graph, _snapshot
+
+__all__ = ["DenseExecutor"]
+
+_runtime_factory: type[Any] | None = None
+if os.environ.get("MLX_ENABLE_TF32") == "0":
+    try:
+        from xarray_sql._dense_mlx import _Runtime
+    except ImportError:
+        pass
+    else:
+        _runtime_factory = _Runtime
+
+
+class _Compiled:
+    def __init__(self, owner, frame):
+        self.schema = frame.schema()
+        self.program = DenseProgram(frame.logical_plan(), self.schema)
+        self.epoch = owner._epoch
+        owner._counts["dense_programs"] += 1
+
+    def run(self, owner):
+        started = time.perf_counter()
+        result = self.program.run(owner._read)
+        owner._counts["dense_executions"] += 1
+        owner._times["dense_lowering_seconds"] += time.perf_counter() - started
+        return result
+
+    def store(self, owner, name, table):
+        owner._store[name] = table
+        owner._stub(name, self.schema)
+
+
+class _DenseBackend(_Backend):
+    def compile(self, plan):
+        return _Compiled(self.owner, _ddx._dataframe(self.ctx, plan))
+
+    def materialize_compiled(self, name, plan, compiled):
+        compiled.store(self.owner, name, compiled.run(self.owner))
+
+    def returns_rows_compiled(self, plan, compiled):
+        table = compiled.run(self.owner)
+        # Checks may have arbitrary relational shapes; materialize their
+        # actual values, never infer emptiness from schema-only providers.
+        return table.arrow(compiled.schema).num_rows > 0
+
+    def drop_table(self, name):
+        self.owner._store.pop(name, None)
+        self.owner._stubs.discard(name)
+        super().drop_table(name)
+
+
+class _DenseGraph(_Graph):
+    def run(self, stack=(), *, rebind=False):
+        owner, aliases = self.owner, {}
+        for name in sorted(self.tables & owner._views.keys()):
+            if name in stack:
+                raise ValueError("cyclic virtual views cannot be staged")
+            graph = owner._view_graph(name)
+            table = graph.run((*stack, name))
+            graph.compiled[""].store(owner, graph.names[""], table)
+            aliases[name] = f'SELECT * FROM "{graph.names[""]}"'
+        for name in self.order:
+            compiled = self.compiled.get(name)
+            if compiled is None or compiled.epoch != owner._epoch:
+                compiled = self.compiled[name] = _Compiled(
+                    owner,
+                    SessionContext.sql(owner.ctx, self._sql(name, aliases)),
+                )
+            table = compiled.run(owner)
+            if name:
+                compiled.store(owner, self.names[name], table)
+        return table
+
+
+class DenseExecutor(SharedExecutor):
+    """Execute supported SQL regions using compact arrays across ddx stages.
+
+    Input Datasets are read afresh on every call, including in-place
+    writes. All final Arrow tables are realized before returning, and no
+    parameter array is mutated during execution. ``fallback=False`` is a
+    diagnostic mode that reports unsupported regions immediately.
+    """
+
+    def __init__(
+        self,
+        ctx,
+        *,
+        precision="float64",
+        fallback=True,
+        max_array_bytes=1 << 30,
+        max_cached_bytes=8 << 30,
+        **options,
+    ):
+        if precision not in {"float64", "float32"}:
+            raise ValueError("precision must be float64 or float32")
+        if precision == "float32" and os.environ.get("MLX_ENABLE_TF32") != "0":
+            raise ValueError("launch Python with MLX_ENABLE_TF32=0")
+        if precision == "float32" and _runtime_factory is None:
+            raise ImportError("MLX is required for float32 GPU execution")
+        if max_array_bytes <= 0:
+            raise ValueError("max_array_bytes must be positive")
+        if max_cached_bytes <= 0:
+            raise ValueError("max_cached_bytes must be positive")
+        super().__init__(ctx, **options)
+        self._max_array_bytes = max_array_bytes
+        self._max_cached_bytes = max_cached_bytes
+        self._runtime = (
+            _runtime_factory()
+            if precision == "float32" and _runtime_factory is not None
+            else None
+        )
+        self._backend = _DenseBackend(self)
+        self._store = {}
+        self._stubs = set()
+        self._fallback = SharedExecutor(ctx, **options) if fallback else None
+
+    def register_view(self, name, sql):
+        super().register_view(name, sql)
+        if self._fallback is not None:
+            self._fallback.register_view(name, sql)
+
+    def _view_graph(self, name):
+        if name not in self._graphs:
+            self._graphs[name] = _DenseGraph(self, self._views[name])
+        return self._graphs[name]
+
+    def _stub(self, name, schema):
+        if name not in self._stubs:
+            empty = pa.RecordBatch.from_arrays(
+                [pa.array([], type=f.type) for f in schema], schema=schema
+            )
+            self.ctx.deregister_table(name)
+            self.ctx.register_record_batches(name, [[empty]])
+            self._stubs.add(name)
+
+    def _read(self, name):
+        if name not in self._store:
+            if name not in self.ctx._registered_datasets:
+                raise _Decline(
+                    f"source {name} has no rectangular Dataset contract"
+                )
+            self._store[name] = _read_dataset(
+                name, self.ctx._registered_datasets[name]
+            )
+        return self._store[name]
+
+    def _invoke(self, execute, fallback, tables):
+        with self._lock:
+            started = time.perf_counter()
+            self.ctx._write(tables)
+            self._refresh()
+            self._store.clear()
+            token = _EXECUTION.set(self._runtime)
+            limit_token = _MAX_BYTES.set(self._max_array_bytes)
+            cache = _ComputedCache(self._max_cached_bytes)
+            cache_token = _CACHE.set(cache)
+            try:
+                overridden = (
+                    self.ctx._custom_udfs - self.ctx._trusted_math_udfs
+                ) & (
+                    set(_OPS)
+                    | {
+                        "coalesce",
+                        "nullif",
+                        "sum",
+                        "avg",
+                        "min",
+                        "max",
+                        "count",
+                    }
+                )
+                if overridden:
+                    raise _Decline(
+                        "custom scalar functions require SQL execution"
+                    )
+                return execute()
+            except _Decline as error:
+                self._counts["dense_fallbacks"] += 1
+                self._last_decline = str(error)
+                if fallback is None:
+                    raise
+                # Bind current function definitions when restarting in SQL.
+                self._fallback._epoch += 1
+                self._fallback._backend._catalog = []
+                return fallback()
+            finally:
+                self._store.clear()
+                _EXECUTION.reset(token)
+                _MAX_BYTES.reset(limit_token)
+                _CACHE.reset(cache_token)
+                self._counts["dense_cache_evictions"] += cache.evictions
+                self._counts["dense_peak_cached_bytes"] = max(
+                    self._counts["dense_peak_cached_bytes"], cache.peak
+                )
+                self._times["dense_call_seconds"] += (
+                    time.perf_counter() - started
+                )
+
+    def prepare(self, sql):
+        if calls_ad(sql):
+            run = self.prepare_all([sql])
+            return lambda **tables: run(**tables)[0]
+        graph = _DenseGraph(self, sql)
+        fallback = (
+            self._fallback.prepare(sql) if self._fallback is not None else None
+        )
+
+        def execute():
+            result = graph.run()
+            return _snapshot(result.arrow(graph.compiled[""].schema))
+
+        return lambda **tables: self._invoke(execute, fallback, tables)
+
+    def prepare_all(self, statements):
+        prepared = _prepare_ad(self._backend, statements)
+        finals = {}
+        fallback = (
+            self._fallback.prepare_all(statements)
+            if self._fallback is not None
+            else None
+        )
+
+        def execute():
+            with prepared.run() as rewritten:
+                results = []
+                for sql in rewritten.statements:
+                    if sql not in finals or finals[sql].epoch != self._epoch:
+                        finals[sql] = _Compiled(
+                            self, SessionContext.sql(self.ctx, sql)
+                        )
+                    compiled = finals[sql]
+                    results.append(
+                        _snapshot(compiled.run(self).arrow(compiled.schema))
+                    )
+                return results
+
+        return lambda **tables: self._invoke(execute, fallback, tables)
+
+    def statistics(self):
+        return (
+            super().statistics()
+            | (dict(self._runtime.counts) if self._runtime is not None else {})
+            | {"dense_last_decline": getattr(self, "_last_decline", None)}
+        )

From 7342ed32ac8c827f64ecb0515f45980166fd1333 Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:34:38 -0800
Subject: [PATCH 5/9] Pin the EinFold compact-array feature commit

Keep the compact executor build reproducible without copying native or array-program sources into this repository.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 accelerators/einfold_mlx/build.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/accelerators/einfold_mlx/build.py b/accelerators/einfold_mlx/build.py
index 7991034f..ba80666e 100644
--- a/accelerators/einfold_mlx/build.py
+++ b/accelerators/einfold_mlx/build.py
@@ -17,7 +17,7 @@
 import tempfile
 from pathlib import Path
 
-_REVISION = "0b2f6db892c3cbf9df2c3d775fbcf4c5654df093"
+_REVISION = "feaca83e4cc1c14c7cbb5f5acc734133005ad57f"
 
 
 def main():

From 0bb410f8ad498d3562ed5c294a573289d5e173bd Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:34:39 -0800
Subject: [PATCH 6/9] Keep complete SQL training steps and optimizer state on
 GPU

Own finite rectangular state, reuse ddx primal outputs, and compile gradient, clipping and AdamW programs with explicit mutable inputs and atomic state publication. Add domain-checked GPU indices, checkpointable dropout, real Shakespeare training, independent nanoGPT validation and checkpoint/resume tests.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 .gitignore                                    |   7 +
 README.md                                     |  20 +
 accelerators/ddx_math/README.md               |  44 ++
 accelerators/ddx_math/build.py                |  90 +++
 accelerators/einfold_mlx/README.md            | 361 +++++++++
 .../einfold_mlx/dense-execution-plan.md       | 300 ++++++++
 .../einfold_mlx/resident-training-plan.md     | 132 ++++
 benchmarks/ml/03_shakespeare.py               | 714 ++++++++++++++++++
 benchmarks/ml/README.md                       | 156 +++-
 benchmarks/ml/_shakespeare_torch.py           | 214 ++++++
 tests/test_math.py                            |  99 +++
 tests/test_shakespeare.py                     | 121 +++
 tests/test_training.py                        | 171 +++++
 xarray_sql/math.py                            |  40 +
 xarray_sql/training.py                        | 702 +++++++++++++++++
 15 files changed, 3167 insertions(+), 4 deletions(-)
 create mode 100644 accelerators/ddx_math/README.md
 create mode 100644 accelerators/ddx_math/build.py
 create mode 100644 accelerators/einfold_mlx/README.md
 create mode 100644 accelerators/einfold_mlx/dense-execution-plan.md
 create mode 100644 accelerators/einfold_mlx/resident-training-plan.md
 create mode 100644 benchmarks/ml/03_shakespeare.py
 create mode 100644 benchmarks/ml/_shakespeare_torch.py
 create mode 100644 tests/test_math.py
 create mode 100644 tests/test_shakespeare.py
 create mode 100644 tests/test_training.py
 create mode 100644 xarray_sql/math.py
 create mode 100644 xarray_sql/training.py

diff --git a/.gitignore b/.gitignore
index 8999c64c..dab5339f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,3 +17,10 @@ CHANGELOG.md
 .direnv
 # Pyodide wheels copied in to test the browser demo locally
 docs/demo/*.whl
+# SQL training checkpoints; retain the accompanying text metrics/logs.
+out-shakespeare-sql/
+benchmarks/apple/shakespeare-matched-*/last.npz
+benchmarks/apple/shakespeare-matched-*/best.npz
+benchmarks/apple/shakespeare-resident-*/last.npz
+benchmarks/apple/shakespeare-compiled-*/last.npz
+benchmarks/apple/shakespeare-compiled-*/best.npz
diff --git a/README.md b/README.md
index 42fe05a7..26c5b819 100644
--- a/README.md
+++ b/README.md
@@ -289,6 +289,26 @@ I want to give a special thanks to the following folks and institutions:
   opinion.
 
 
+## Apple silicon tensor acceleration
+
+The optional [Apple silicon integration](accelerators/einfold_mlx/README.md)
+connects [EinFold](https://github.com/xqlsystems/einfold), a tensor optimizer
+for SQL, to [MLX](https://github.com/ml-explore/mlx), Apple's array framework.
+It keeps SQL planning and gradients from
+[ddx](https://github.com/xqlsystems/ddx), which differentiates SQL queries.
+Float64 CPU execution is the default;
+float32 GPU execution requires explicit opt-in. See the
+[local M5 Ultra measurements](benchmarks/apple/README.md) for end-to-end
+results, numerical errors, dispatch thresholds and limitations.
+The [shared-CTE nanoGPT comparison](benchmarks/apple/nanogpt-shared.md)
+combines cached shared SQL stages with the same CPU/GPU kernels and checks
+larger models against JAX.
+The experimental [compact execution path](benchmarks/apple/dense-execution.md)
+retains arrays across forward and backward stages, reuses plans, and
+executes supported operators through MLX. Its report includes checked
+five-step GPT-2-size training, memory bounds and remaining native-framework
+performance gaps.
+
 ## License
 
 ```
diff --git a/accelerators/ddx_math/README.md b/accelerators/ddx_math/README.md
new file mode 100644
index 00000000..ec0eaa55
--- /dev/null
+++ b/accelerators/ddx_math/README.md
@@ -0,0 +1,44 @@
+# Exact GELU for SQL training
+
+[ddx](https://github.com/xqlsystems/ddx) differentiates SQL query plans.
+Its released scalar rules do not include the error function, `erf`, used
+by [nanoGPT](https://github.com/karpathy/nanoGPT)'s exact GELU activation.
+The pinned feature commit adds `erf'(x) = 2/sqrt(pi) * exp(-x*x)` to the scalar
+registry and retains `erf` in derivative plans. The rule composes with
+the chain rule and higher derivatives; it does not substitute a model's
+handwritten backward pass.
+
+The builder uses the [ddx feature fork](https://github.com/cromwellian/ddx), pinned to commit
+`3cd5f2711221e94295f47aa6f1270f712c2edfab`, based on release commit
+`af71683f6678bb093e600f948c212490a4d8b74b`. Historical benchmark wheels used
+`ddxdb==0.4.6+xqlmath.3`; the contribution retains upstream version 0.4.6.
+The feature commit also exposes
+`Prepared.values`: metadata for the primal outputs computed by prepared
+derivative programs. These outputs can be read during a run before its
+temporary tables are released. This lets the GPU training session reuse
+the actual loss from backward preparation instead of executing a separate
+forward query. It does not alter existing derivative rules.
+
+Engines must
+register an `erf(DOUBLE)` implementation; `xarray_sql.math.register_erf`
+does this for DataFusion and marks its equivalent array primitive as safe
+for EinFold CPU and MLX GPU execution. Replacing the SQL UDF removes that
+permission. This contribution does not add an erf implementation to DuckDB.
+
+From the xarray-sql directory, after installing the normal ML dependencies
+and the [compact EinFold executor](../einfold_mlx/README.md):
+
+```bash
+source .venv/bin/activate
+uv pip install cmake maturin
+python accelerators/ddx_math/build.py
+```
+
+Requires git, a Rust toolchain, CMake and uv on PATH. The builder clones
+an isolated checkout, validates the pinned feature revision, builds its
+locked wheel and installs it in the active virtual environment. No patch
+is applied. Use `--source /path/to/ddx` for a clean checkout already at the pin. `uv sync` may
+replace locally built wheels; rerun the optional builders if necessary.
+
+The public SQL numerical checks are in `tests/test_math.py`. The ddx contribution
+also adds a Rust check of the chain rule and second derivative.
diff --git a/accelerators/ddx_math/build.py b/accelerators/ddx_math/build.py
new file mode 100644
index 00000000..a77c4443
--- /dev/null
+++ b/accelerators/ddx_math/build.py
@@ -0,0 +1,90 @@
+"""Build ddx's exact error-function derivative in an isolated checkout.
+
+ddx (https://github.com/xqlsystems/ddx) differentiates SQL. The pinned
+feature commit adds erf's analytic derivative, including higher derivatives, and
+exposes prepared primal output metadata for loss reuse.
+Requires git, cargo, CMake, maturin and uv on PATH.
+"""
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+_REVISION = "3cd5f2711221e94295f47aa6f1270f712c2edfab"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--source", type=Path, help="reuse a clean pinned checkout"
+    )
+    args = parser.parse_args()
+    if sys.prefix == sys.base_prefix:
+        parser.error("use the repository virtual environment's Python")
+    for command in ("git", "cargo", "cmake", "uv"):
+        if shutil.which(command) is None:
+            parser.error(f"{command} must be on PATH")
+    with tempfile.TemporaryDirectory(prefix="xql-ddx-math-") as temporary:
+        source = args.source or Path(temporary) / "source"
+        if args.source is None:
+            subprocess.run(
+                [
+                    "git",
+                    "clone",
+                    "https://github.com/cromwellian/ddx.git",
+                    str(source),
+                ],
+                check=True,
+            )
+            subprocess.run(
+                ["git", "checkout", _REVISION], cwd=source, check=True
+            )
+        if (
+            subprocess.check_output(
+                ["git", "rev-parse", "HEAD"], cwd=source, text=True
+            ).strip()
+            != _REVISION
+        ):
+            parser.error(f"expected ddx revision {_REVISION}")
+        if subprocess.check_output(
+            ["git", "status", "--porcelain"], cwd=source
+        ):
+            parser.error("source checkout must be clean")
+        wheels = Path(temporary) / "wheels"
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "maturin",
+                "build",
+                "--release",
+                "--locked",
+                "--out",
+                str(wheels),
+            ],
+            cwd=source / "python/ddxdb",
+            env=os.environ | {"PYO3_PYTHON": sys.executable},
+            check=True,
+        )
+        (wheel,) = wheels.glob("*.whl")
+        subprocess.run(
+            [
+                "uv",
+                "pip",
+                "install",
+                "--python",
+                sys.executable,
+                "--no-deps",
+                "--reinstall",
+                str(wheel),
+            ],
+            check=True,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/accelerators/einfold_mlx/README.md b/accelerators/einfold_mlx/README.md
new file mode 100644
index 00000000..d82e001c
--- /dev/null
+++ b/accelerators/einfold_mlx/README.md
@@ -0,0 +1,361 @@
+# EinFold kernel bridge for Apple silicon
+
+[EinFold](https://github.com/xqlsystems/einfold) recognizes tensor operations
+in Apache DataFusion SQL plans. [MLX](https://github.com/ml-explore/mlx) is
+Apple's array framework for Apple silicon. This implements Option B:
+
+```text
+SQL → DataFusion logical plan → ddx rewrite, when requested
+    → EinFold detection and SQL input validation
+    → float64 CPU matrix multiplication or float32 MLX GPU kernels
+    → EinFold row/domain reconstruction → DataFusion → Arrow/xarray
+```
+
+[ddx](https://github.com/xqlsystems/ddx), the SQL automatic differentiation
+engine, constructs every gradient, VJP and JVP. Logical plans remain normal
+DataFusion plans until physical planning, so ddx can serialize them through
+Substrait. MLX only evaluates arithmetic inside those plans.
+
+The released EinFold Python API has no external executor hook. The optional
+builder pins the feature commit `feaca83e4cc1c14c7cbb5f5acc734133005ad57f` in the [EinFold feature fork](https://github.com/cromwellian/einfold),
+which contains the native bridge and compact array module. Its tested base
+is `39b25210d3b39027ef91584383fb114dfa8a4a48` (Python 0.2.0). Historical
+benchmark wheels were labeled `0.2.0+xqlmlx.6`; the contributed source retains
+version 0.2.0. This feature is not in the ordinary PyPI wheel yet. The
+xarray-sql Rust extension supplies SQL parsing and verified ordered Arrow
+providers used by shared execution.
+
+## Owned state and complete training graphs
+
+`xarray_sql.training.TrainingSession` adds an explicit ownership boundary
+for complete, finite, nonnullable rectangular state tables. It imports
+weights and optimizer moments once, then retains MLX arrays between SQL
+updates. Ordinary `DenseExecutor` queries retain their existing live-input
+and Arrow snapshot behavior. Use a dedicated Context for a session and
+read its current state through `session.query` or `session.export`.
+
+`session.prepare_step` receives the original SQL gradient statements,
+ordinary SQL intermediate/optimizer statements and a mapping from their
+outputs to persistent state tables. It evaluates gradients, clipping and
+all parameter updates in one invocation. The optional
+[ddx math/primal build](../ddx_math/README.md) exposes the already computed
+primal loss. No independent loss-forward call is needed.
+
+With `compile=True`, the session traces these SQL/ddx array programs into
+one cached `mx.compile` function, with mutable numerical state passed as
+explicit array arguments. Schemas, coordinate layouts, view definitions
+and session settings guard reuse. Structural checks still run or must be
+proven from the coordinate/count facts; a check that requires an
+unsupported dynamic host decision declines compilation. Device finite,
+NULL and complete-domain checks return a status scalar before all new
+state tables are published together. Failed updates retain the previous
+state and RNG counter.
+
+The pinned EinFold integration supplies device tensor leaves, finite integer
+domain metadata, GPU index joins and embedding scatter grouping. Declare
+bounded integer inputs with `session.index_domains` before binding them.
+Actual input values are checked against those domains. Numerical index
+arithmetic requires an explicit floating cast; unsupported integer/domain
+behavior declines instead of silently changing SQL arithmetic. Schema or
+coordinate changes require a new session. Explicit `bind` replacements
+copy caller inputs and preserve ownership; exports are independent Arrow
+snapshots. Numerical host conversion during compilation is refused.
+
+`session.dropout` generates masks from an explicit seed and step key on
+GPU. Save the seed and next step with the checkpoint. Complete training
+results can retain masks and raw gradients for correctness inspection;
+these exports are outside benchmark timers. Ordinary steps read only the
+validation and loss scalars. Full state is exported for checkpoints.
+
+From the xarray-sql directory, the matching Shakespeare recipe uses:
+
+```sh
+MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/03_shakespeare.py \
+  --compile-training --gpu-dropout --steps 100
+```
+
+Omit `--steps` for the recipe's full training run. `--resident` exercises
+the intermediate implementation with device state and shared primal loss,
+without whole-step compilation. `--compile-training` without GPU dropout
+keeps NumPy-generated masks as explicit inputs, useful for comparison with
+the earlier deterministic trajectory. CPU resident execution uses
+`--resident --backend dense_cpu --cache-gib 32`.
+
+## Compact execution across forward and backward stages
+
+`xarray_sql.dense.DenseExecutor` is an opt-in implementation of
+SQL → DataFusion logical plans → EinFold array programs → NumPy/MLX.
+It retains shared CTE and ddx stage values as private arrays with explicit
+coordinate axes. Numerical joins become gathers or contractions, window
+reductions broadcast, and reduced UNION branches accumulate gradients.
+Coordinates expand into scalar Arrow columns at public output boundaries.
+The executor orders compact axes before expansion where their exact domains
+prove the requested SQL ordering.
+
+```python
+from einfold.datafusion import EinFold
+from xarray_sql import Context
+from xarray_sql.dense import DenseExecutor
+
+ctx = Context(threads=30).with_extensions(EinFold())
+# Register input and parameter Datasets on ctx.
+executor = DenseExecutor(ctx)  # float64 CPU by default
+# Explicitly permit float32 GPU arithmetic, with MLX_ENABLE_TF32=0 at launch:
+# executor = DenseExecutor(ctx, precision="float32")
+executor.register_view("model", model_select_sql)
+loss = executor.prepare(loss_select_sql)
+step = executor.prepare_all(gradient_update_sqls)
+loss_rows = loss()
+updated_parameter_rows = step()
+print(executor.statistics())
+```
+
+ddx differentiates the original SQL. This path does not use MLX automatic
+differentiation or replace a model with framework code. Logical programs
+and compiled elementwise functions are reused across steps; compiled
+functions receive every numerical leaf as an argument. Input NumPy buffers
+are imported afresh each call, preserving in-place writes and registration
+changes. Private schema-only providers serve ddx's named reads; they are
+never scanned to produce numerical results. Validation checks evaluate
+actual rows. Public outputs are independent Arrow snapshots, and input
+parameters change only when the caller writes an update.
+
+The compact source contract is a rectangular Dataset with unique integer
+or string coordinates and finite float64 numerical buffers. The initial
+implementation supports a closed set of arithmetic, CASE/NULL operations,
+gathers, grouped reductions, full-partition windows, axis joins and domain
+left joins. Unsupported layouts or operations restart the call through
+the existing shared SQL executor before any parameter is changed.
+`fallback=False` reports unsupported cases directly and is used by the
+benchmarks. Custom scalar-function overrides stay in SQL. GPU precision is
+explicit; it retains the existing `3e-5` benchmark tolerance and the
+float32 AdamW limitations documented in the measurements.
+
+The `.5` integration also retains multiplicative factors through aliases
+and zero-valued CASE branches when reducing a conditional product. This
+keeps attention backward masks inside contractions instead of expanding a
+batch × heads × context × context × channels temporary. Exceptional
+arithmetic retains the normal SQL fallback. Exact GELU is available with
+the optional [ddx erf build](../ddx_math/README.md) and
+`xarray_sql.math.register_erf(ctx)`; arbitrary UDF overrides remain in SQL.
+The [Shakespeare trainer](../../benchmarks/ml/03_shakespeare.py) uses
+these features for dropout and AdamW at nanoGPT's character-model size.
+
+GPU intermediates remain MLX arrays through their supported operations.
+Compiled pointwise regions use `mx.compile`; contractions use MLX matrix
+operations and embedding gradients use scatter reductions. Exact counts
+use integer arithmetic, with bounded int32 scatter where Metal requires
+it. GPU readbacks occur at SQL validation/output boundaries. Caches retain
+program structure across calls; numerical graphs are released after each
+call. This API still imports and exports the full parameter state each
+training call. An API that owns persistent GPU training state is a future
+extension with a different ownership and synchronization contract.
+
+The computed-value cache defaults to 8 GiB per invocation; adjust it with
+`max_cached_bytes`. Source buffers remain available throughout the call;
+computed values may be evicted and recomputed. `max_array_bytes` defaults
+to a conservative 1 GiB estimate for a numerical temporary. These are
+array/cache bounds, not an RSS limit: source tables, public Arrow results,
+working allocations and driver memory are additional. The benchmark
+runner enforces its own process limit and accepts `--dense-cache-gib`.
+Exact coordinate and nonnullable-domain facts keep row and NULL masks
+compact, including UNION branches whose domains differ on one axis.
+
+The [dense execution report](../../benchmarks/apple/dense-execution.md)
+contains measured speedups, memory, copy counters, failed attempts and
+remaining overheads. The
+[research and execution plan](dense-execution-plan.md) records the source
+research and implementation sequence.
+
+## Install and run
+
+Requires Apple silicon, a working Metal GPU, Python 3.10+, Rust 1.94+,
+Git, uv and maturin. In the xarray-sql checkout:
+
+```sh
+uv sync --extra mlx --extra ml
+.venv/bin/python accelerators/einfold_mlx/build.py
+MLX_ENABLE_TF32=0 .venv/bin/python your_program.py
+```
+
+The builder checks out the pinned feature commit into a temporary directory,
+builds its locked wheel, and installs only into the interpreter's environment.
+No patch is applied and no EinFold source is duplicated in this repository.
+A later `uv sync` can replace the custom wheel; rerun the builder afterward.
+Use `--source /path/to/einfold` to build a clean checkout already at the pin.
+For supplied benchmarks, `.venv/bin/python -m benchmarks.apple.benchmark`
+runs this checkout directly. Directly running an ML demo needs
+`PYTHONPATH="$PWD"` unless xarray-sql is installed as editable.
+
+The pinned EinFold wheel includes `einfold.dense_program`, used by the
+optional `DenseExecutor` to retain compact arrays across SQL stages.
+
+```python
+from xarray_sql import Context, ad
+from xarray_sql.mlx import MLX
+
+accelerator = MLX(precision="float32")
+ctx = Context(threads=30).with_extensions(accelerator)
+ctx.from_dataset("x", dataset)
+rows = ctx.sql(statement).to_arrow_table()
+print(accelerator.statistics())
+
+# One SQL gradient query can use the same preparation API as forward SQL.
+single_step = ctx.prepare(gradient_statement)
+rows = single_step(x=new_dataset_values)
+
+# Several parameter updates can share one loss and its ddx programs.
+step = ad.prepare(ctx, [gradient_statement, update_statement])
+gradient_rows, updated_rows = step(x=new_dataset_values)
+```
+
+`Context.with_extensions()` returns another xarray-aware handle, sharing
+Dataset metadata and prepared-query invalidation with the original handle.
+Always use the returned handle, which carries the extension's codecs.
+`threads=` controls DataFusion's target parallel partitions. The default is
+all 30 CPU cores on the measured M5 Ultra; 60 partitions modestly helped
+plain DataFusion attention but slightly slowed EinFold in the confirmation
+run. The [overhead report](../../benchmarks/apple/README.md#training-overheads-and-cpu-parallelism)
+contains both results.
+
+## Numerical and SQL contracts
+
+- `MLX()` defaults to **float64 CPU** matrix multiplication through NumPy;
+  attention keeps EinFold's original CPU implementation. No implicit
+  conversion of double precision scientific data is enabled by default.
+- `MLX(precision="float32")` explicitly permits single precision GPU
+  arithmetic, including intermediate values used by ddx. Outputs keep the
+  SQL Arrow dtype. **It does not preserve float64 rounding or error bounds.**
+- Start Python with `MLX_ENABLE_TF32=0`. On M5, MLX otherwise uses arithmetic
+  with less precision than float32 for matrix operations. The extension
+  requires this setting and does not mutate process-wide MLX settings.
+  See [MLX's numerical precision documentation](https://ml-explore.github.io/mlx/build/html/usage/precision.html).
+- EinFold retains key equality, bag semantics, NULLs, row existence,
+  grouping, domain defaults, and output mask reconstruction. Duplicate,
+  sparse, non-positional and unsupported inputs keep EinFold/DataFusion
+  fallback behavior. Non-finite factors stay on the CPU. Very large or tiny
+  finite factors also decline GPU execution.
+- GPU attention is limited to complete finite inputs, product scores,
+  shifted exponentials and normalized weights. Unshifted exponentials,
+  general score expressions and other pipelines use existing CPU kernels.
+  Key intervals come from EinFold's normalized mask conditions. Masked
+  infinities never enter GPU arithmetic.
+- There is no input-value cache: repeated queries and in-place parameter
+  writes read current values. Backend ownership is session-scoped; ordinary
+  DataFusion and EinFold contexts can coexist with accelerated contexts.
+- Gradient definitions and typed ddx refusals are unchanged. `ad.prepare`
+  reuses ddx programs while checking catalog/schema/settings changes;
+  it materializes final results so later calls cannot overwrite them.
+  It intentionally uses ddx's regular preparation, which supports table
+  replacement, rather than experimental in-place physical-plan reuse.
+
+## Dispatch and memory
+
+The native side declines small requests before copying buffers into Python
+(default `min_ops=2**20` scalar products). In float32 mode, GPU dispatch
+starts at `min_gpu_ops=2**33`; smaller products use float64 CPU matrix
+multiplication and smaller attention uses EinFold's CPU tiles. This avoids
+GPU dispatch regressions on small workloads. Both limits can be adjusted;
+`min_ops=0, min_gpu_ops=0` forces eligible computations through the GPU for
+correctness testing or threshold experiments.
+
+Attention uses at most 1024 queries per tile, with a smaller tile when the
+backend's working-memory estimate requires it. Masks and temporary score
+storage grow with the number of keys and tile width, rather than the full
+number of query/key pairs. Each tile trims keys outside the union of its
+allowed intervals. Rectangular and causal intervals use MLX's mask-free or
+causal kernels when their ranges match exactly; other masks remain explicit. Backend buffers are bounded by a configurable
+512 MiB estimate and reserved through EinFold's DataFusion memory budget.
+This estimate is not a hard limit on the GPU driver's allocator or its
+kernel cache. All GPU results synchronize before Arrow conversion.
+
+`statistics()` counts completed `cpu_matmul`, `gpu_matmul`, and
+`gpu_attention` kernels, plus requests declined inside Python. Requests
+rejected by native thresholds are not counted. Planning a GPU-capable
+operator alone does not increment a completion counter.
+
+## Validation
+
+See [the measured report](../../benchmarks/apple/README.md), the standalone
+benchmark, raw timings and CPU profile in `benchmarks/apple`.
+
+```sh
+MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python -m pytest -m 'not integration'
+MLX_ENABLE_TF32=0 .venv/bin/python -m pytest tests/test_mlx.py
+```
+
+The native patch also passes EinFold's kernel and randomized invariant
+suite, except for the existing plain-DataFusion `nan_compares_by_its_sign`
+assumption on this ARM Mac. That test expects `inf - inf` to produce a
+negative NaN. Here plain DataFusion produces a positive NaN, without any
+extension installed. The test remains unchanged in the patch and was
+documented in the measured report; it is not silently removed.
+
+## Shared CTE execution
+
+The `0.2.0+xqlmlx.3` patch adds safe prepared-plan resets and multi-key
+attention detection. Batch and head keys must match in the scores, values,
+softmax partition and normalizer. The kernel validates complete positional
+domains before treating those groups as independent attention problems.
+
+`xarray_sql.shared.SharedExecutor` materializes reachable CTEs read more than
+once and caches each stage's physical plan. It writes new stage values into
+owned Arrow buffers. A changed coordinate domain, row count or schema
+invalidates plans that read the replaced buffer. Dynamic filter pushdown is
+disabled on this explicit executor's context, because its cached filters
+would otherwise retain values from an earlier training step.
+
+This reset also runs when DataFusion's FFI wrapper replaces children:
+DataFusion 55 does not forward a separate reset callback through that
+wrapper. Ordinary `Context.prepare` conservatively retains only the
+optimized logical plan when dynamic filters are enabled.
+
+```python
+from xarray_sql import Context
+from xarray_sql.mlx import MLX
+from xarray_sql.shared import SharedExecutor
+
+accelerator = MLX(precision="float32", min_ops=0, min_gpu_ops=0)
+ctx = Context().with_extensions(accelerator)
+# Register parameter/input Datasets on ctx first.
+executor = SharedExecutor(ctx, inline_ctes=["scores", "weights", "norm"])
+executor.register_view("model", model_select_sql)
+loss = executor.prepare(loss_select_sql)
+step = executor.prepare_all(gradient_update_sqls)
+```
+
+Virtual views retain their original SQL. ddx receives those definitions,
+not the forward executor's materialized tables, and builds the backward
+pass before its own stages execute. Numerical work in both paths passes
+through DataFusion and EinFold before the MLX callback. Returned Arrow
+tables are independent snapshots.
+
+DataFusion 55's Substrait consumer cannot apply a projected read to an
+expanded view. When a ddx stage reads a deferred view, the executor binds
+the same program through ddx's SQL writer instead. It retains the original
+SQL differentiation and avoids substituting forward materializations into
+the derivative's lineage.
+
+Join-only ddx stages remain logical views until a consumer reduces their
+rows, avoiding storage of the full contraction join. Very large cotangent
+sums (at least 32 UNION branches) materialize their reductions sequentially
+to limit concurrent join builds. Compound reduction stages rebind transient
+ddx providers on each call. They retain output storage, while ordinary
+forward and backward stages also retain physical plans.
+
+An exact, bounded coordinate cache shares immutable integer/string buffers;
+floating point values always have separate writable storage. Every cache
+hit checks all values and NULLs. The default budget is 256 MiB, adjustable
+with `max_key_cache_bytes`. Canonically ordered stage tables expose verified
+ordering to DataFusion without an additional sort. Coordinate/domain/schema
+changes invalidate dependent plans. See the
+[shared nanoGPT measurements](../../benchmarks/apple/nanogpt-shared.md)
+for measured benefits and remaining large-model memory limits.
+
+This API explicitly supports deterministic numerical SELECTs with flat,
+nonrecursive WITH clauses. It refuses unknown functions and nested WITH
+scopes. Callers identify CTEs that should remain inside a fused region with
+`inline_ctes`. Stage boundaries still exchange Arrow data: this does not
+implement a GPU-resident training graph or replace SQL derivatives with
+MLX autodiff. The native xarray-sql extension must also be rebuilt for its
+SQL-parser helper (`maturin develop --release` in the checkout).
diff --git a/accelerators/einfold_mlx/dense-execution-plan.md b/accelerators/einfold_mlx/dense-execution-plan.md
new file mode 100644
index 00000000..d56412f1
--- /dev/null
+++ b/accelerators/einfold_mlx/dense-execution-plan.md
@@ -0,0 +1,300 @@
+# Dense execution of SQL training programs
+
+Implementation and validation results are available in the
+[dense execution report](../../benchmarks/apple/dense-execution.md).
+The local executor now retains compact forward/backward arrays, reuses
+structural plans, runs the supported program on MLX, preserves the original
+ddx differentiation, and has completed checked five-step GPU training at
+context 1024. Default APIs still return Arrow snapshots. An explicit API
+for persistent GPU-owned training state remains future work; long-context
+CPU cache sensitivity and the dense AdamW refusal are reported as limits.
+
+Research and execution plan, November 1, 2026. The requested execution
+threshold is 06:00 America/Los_Angeles (14:00 UTC); research began before
+it, and implementation may proceed now that this time has passed.
+
+## Evidence and objective
+
+The current SQL → DataFusion → EinFold → MLX bridge accelerates individual
+matrix and attention kernels, but retains relational intermediates. The
+123,563,520-parameter model (12 layers, width 768, 12 heads, vocabulary
+50257, batch 1, context 16) takes 163.65 s per warm SQL GPU step. Across
+five steps, counters record 131 s planning, 169 s stage storage and 504 s
+execution, with 7,148 physical plans built. Native JAX CPU takes 0.334 s
+eager, 0.394 s with separate loss/backward calls and 0.136 s compiled;
+upstream nanoGPT's PyTorch model takes 0.095 s CPU float64 and 0.024 s MPS
+float32. All parameter updates were checked. These are short-context
+training results; the existing context-1024 comparison is forward-only.
+
+See [SQL measurements](../../benchmarks/apple/nanogpt-shared.md),
+[native measurements](../../benchmarks/apple/nanogpt-frameworks.md) and
+their linked raw records. They are immutable baselines for this work.
+
+The objective is to execute validated tensor-shaped regions of the actual
+SQL and ddx programs on compact arrays, keeping intermediates dense and
+avoiding per-operator Arrow reconstruction. ddx remains the source of
+derivatives. The implementation must work on plans and axis relationships,
+without recognizing a benchmark's name or substituting a handwritten GPT.
+
+## Research findings
+
+[EinFold](https://github.com/xqlsystems/einfold) is a tensor optimizer for
+SQL. Its [design](https://github.com/xqlsystems/einfold/blob/main/docs/design.md)
+already describes exact shape/layout facts in a side table, separate
+symbolic and numerical passes, and optimization of a whole program of
+plans. Its current DataFusion operators still collect Arrow inputs and
+return relational Arrow rows. This supports adding a dense program path
+without changing the SQL surface, while device-specific execution remains
+in the host integration. A design aspiration is not an implemented API.
+
+[ddx](https://github.com/xqlsystems/ddx) differentiates SQL. Its
+[program representation](https://github.com/xqlsystems/ddx/blob/main/crates/ddx-ad/src/program.rs)
+contains input tables, validation checks, forward steps, backward steps
+and output tables. Each step has a name and an unbound Substrait plan;
+reads can be bound to schemas. The Python backend currently exposes
+per-step `compile`, `materialize_compiled` and existence checks. Those
+hooks permit an incremental dense backend. A whole-program preparation
+hook is a possible later extension, not a precondition for useful progress.
+
+[DataFusion](https://datafusion.apache.org/library-user-guide/extending-operators.html)
+supports custom logical/physical nodes through optimizer rules. A dense
+region can therefore be represented internally while preserving its
+ordinary relational input/output contract. One physical region should
+contain many array operations, rather than returning scalar Arrow rows
+between every pair of operators.
+
+[MLX](https://ml-explore.github.io/mlx/build/html/usage/unified_memory.html)
+is Apple's array framework for unified CPU/GPU memory. The important costs
+here are allocation, packing, conversion and synchronization; CPU/GPU
+crossings on this machine are not necessarily discrete-device PCIe copies.
+[Its conversion documentation](https://ml-explore.github.io/mlx/build/html/usage/numpy.html)
+states that `mx.array` copies, and float64 NumPy input converts to float32.
+Buffer/DLPack interop has ownership and synchronization requirements.
+Unified memory alone does not make our current bridge zero-copy.
+
+[MLX compilation](https://ml-explore.github.io/mlx/build/html/usage/compile.html)
+can eliminate common work and fuse operations. Shape/type/input changes
+can cause recompilation. Build functions over explicit array arguments;
+avoid capturing mutable parameter values in cached closures. Defer
+evaluation to region/output boundaries, with explicit synchronization for
+timing. Execute ddx's derivative expressions; do not use MLX autodiff to
+replace them. Keep `MLX_ENABLE_TF32=0` as required by the
+[precision documentation](https://ml-explore.github.io/mlx/build/html/usage/precision.html).
+
+[Arrow's fixed-shape tensor type](https://arrow.apache.org/docs/format/CanonicalExtensions.html#fixed-shape-tensor)
+is useful for tensor interchange but changes a scalar field into a tensor
+field. It cannot silently replace a public SQL table of scalar values.
+Use a private dense representation and reconstruct the ordinary rows at
+real relational/output boundaries. The
+[Arrow C interface](https://arrow.apache.org/docs/format/CDataInterface.html)
+can share owned Arrow buffers across runtimes, but does not prove tensor
+shape, uniqueness, absence of NULLs, or MLX compatibility.
+
+## Recommended design
+
+Use a program-scoped dense store and a small array-operation graph.
+Preserve existing public SQL/Arrow schemas and a relational fallback.
+
+### Compact values and exact domains
+
+A dense value contains its array/buffer, ordered axes, dtype, shape,
+strides/view mapping and lifetime owner. Store coordinate dictionaries
+once per axis/domain, rather than once per tensor element. Column lineage
+maps SQL fields to axis coordinates, indexed gathers or scalar array
+expressions. Attach exact facts to the plan node where they are valid.
+
+Initial eligibility is deliberately narrow: verified rectangular domains,
+unique keys and supported finite numerical inputs. Arbitrary coordinate
+labels may use exact dictionaries. Unsupported duplicate, NULL, sparse or
+value-dependent domains retain relational execution until their handling
+is proved. Shape is insufficient evidence of completeness or uniqueness.
+
+Cache structural proofs under coordinate/schema epochs. Refresh numerical
+values under separate value epochs. Ordinary mutable NumPy/Arrow inputs
+are imported once per call, so direct external writes remain visible.
+Cross-call value caching is allowed only for explicitly owned, versioned
+state with a defined write path. Do not infer that an untracked external
+array has stayed unchanged because its schema or pointer is unchanged.
+
+### Dense lowering of backward subgraphs
+
+Lower a closed set of verified operations: reads, aliases, projections,
+views, elementwise arithmetic/CASE, broadcasts, gathers, reductions,
+contractions and required domain-fill operations. Use field lineage and
+axis relationships from actual DataFusion/Substrait plans. Avoid SQL text
+regexes and neural-network-specific substitutions.
+
+Recognize shared-axis join plus grouped product as a contraction before
+forming its Cartesian products. For example, the tied embedding gradient
+has the form `X.T @ dlogits`: context 16, width 768 and vocabulary 50257
+represent 617,558,016 possible product contributions, but only a
+38,597,376-element output. Retain that matrix formulation through normalized
+and residual operands. Extend existing EinFold fold/chain detection where
+that removes a dominant missed pattern; use a general dense operation
+graph for the surrounding maps and reductions.
+
+Priority patterns are reconstructed layer-normalized operands, residual
+joins, null-safe key joins over proven nonnullable axes, gradient domain
+left joins with defaults, gather/scatter contributions for embeddings,
+and sums of many reduced gradient contributions. A UNION is a bag union,
+not elementwise addition unless a following reduction and domain proof
+establish that equivalence. Group existence and all-NULL SUM behavior must
+be retained. Masks describe absent rows, with no multiplication of masked
+NaN/Inf values by zero.
+
+### Stable bindings and compiled programs
+
+Bind intermediate reads to stable logical slots rather than new physical
+table providers on each run. Each slot has an immutable structural
+descriptor and a current value generation. The graph cache key includes
+schema, coordinate/domain facts, function semantics/settings, dtype and
+precision mode. Changing only parameter values updates bindings; changed
+domains or semantics invalidate affected regions.
+
+Start through ddx's existing backend hooks: compiled steps may hold dense
+plans and `materialize_compiled` may store a lazy internal value. Named
+read schemas must remain real and fallback scans must materialize current
+rows on demand. Existence checks can use exact domain facts only when
+their result is proved; otherwise evaluate the check or cross a relational
+boundary. Public methods returning Arrow still return complete Arrow
+snapshots.
+
+If per-step hooks prevent adequate fusion, add an optional ddx
+whole-program compilation interface using its existing program data.
+Existing backends keep their current behavior. Preserve validation checks,
+typed refusals and the ordering of externally visible effects.
+
+### CPU first, then resident MLX regions
+
+Implement the dense path on NumPy float64 first to distinguish relational
+lowering problems from GPU precision/device behavior. Then execute the
+same operation graph with MLX arrays in explicitly permitted float32 mode.
+Keep maximal supported regions as arrays; convert at imports, unsupported
+relational boundaries and requested public outputs. Avoid `np.asarray`,
+`.item()` and `mx.eval` inside each individual numerical operator.
+
+The current slice-based native backend API forces packing into f64 vectors
+and returning f64 vectors per call. The new path needs an owned array/graph
+interface with stable handles and views, while retaining the existing
+callback as a fallback. Keep EinFold's IR hardware-neutral and MLX objects
+in the host bridge. Profile before adding custom C++/Metal kernels.
+
+Only after the compatible Arrow-output path is correct, consider an
+explicit resident training-state API. It must state who owns parameters,
+how writes are versioned and when host snapshots are produced. Existing
+`Context.prepare`/`SharedExecutor` APIs must not silently change their
+return types or ignore external writes. Compare resident compute and
+Arrow-materializing API time separately.
+
+### Memory and fallback
+
+Use last-use analysis to release intermediate arrays and accumulate
+gradient contributions as they become available. Bound allocations using
+known shapes, include conversion/output costs and avoid retaining every
+branch's large result. MLX's lazy graph must also be bounded and evaluated
+at deliberate points rather than growing across unlimited training steps.
+
+Failed runtime guards must fall back before committing parameter updates.
+Retain an unchanged input snapshot or otherwise transactional ownership
+until guards pass. An unsupported region is a normal relational boundary,
+not permission to change NULL/duplicate/default semantics. Use float64 CPU
+as the default exact-precision path. Do not weaken tolerances or silently
+alter AdamW to hide the existing float32 parity failure.
+
+## Execution sequence
+
+1. **Expose the actual bottlenecks.** Add opt-in per-stage timings and
+   dense/fallback reasons, row/byte counts, coordinate packing costs,
+   compiled-plan counts and synchronization/conversion counters. Profile
+   one representative backward pass, not repeated five-step large runs.
+   Save top offenders and minimal SQL reproductions. Native opaque nodes
+   require internal metrics; their outer node count is not an operator
+   count comparable with plain DataFusion.
+2. **Build dense storage and lowering on CPU.** Add exact domains, compact
+   arrays and validated operator lowering. First cover ordinary matrix
+   gradients and reconstructed normalized/residual operands. Route those
+   through actual SQL/ddx plans. Keep unsupported operators on fallback.
+   Test different names, shapes and coordinate labels to prevent accidental
+   model-specific behavior.
+3. **Integrate stable staged execution.** Connect dense values to ddx's
+   compiled backend hooks and stable slots. Share intermediates and
+   eliminate per-step replanning for unchanged shapes. Verify repeated
+   updates, domain changes, empty/NULL results and snapshots. Preserve
+   dynamic-filter safeguards and fresh join state for relational regions.
+4. **Execute dense regions with MLX.** Compile supported graphs over
+   explicit parameters, retain arrays between their operators, evaluate
+   at required boundaries and record all copies/casts/synchronizations.
+   Add explicit owned resident state only after the compatible path works.
+5. **Expand supported backward patterns.** Use the measured remaining
+   fallback list to add domain fill, scatter/reduction, nonlinear maps and
+   broader contraction recognition. Recheck memory liveness and accuracy
+   after each coherent change; avoid an unbounded optimizer rewrite search.
+6. **Validate and benchmark.** Run focused public-result tests, then the
+   full local suite and native tests appropriate to changed crates. Start
+   with tiny/213K five-step cases and normalized-gradient reproductions,
+   then 10.7M and 123.6M context-16 five-step comparisons. Attempt larger
+   contexts progressively only after short-context correctness and memory
+   are established. Run hosts serially without builds/profilers during
+   timings. Compare unchanged JAX/PyTorch controls; keep failed/bounded runs
+   out of speedup tables. Include DuckDB controls where bounded runs
+   complete. Write measured speedups and remaining limitations.
+
+These are ordered milestones, not permission to stop after a fast kernel
+microbenchmark. Adapt implementation choices to measured blockers while
+preserving the scope and correctness requirements.
+
+## Completion criteria
+
+- SQL and ddx remain the source of the computation and derivatives.
+- Supported dense backward paths avoid product-sized relational
+  intermediates and per-element coordinate storage between their operators.
+- Steady-state structural planning is independent of training-step count
+  for unchanged domains; value changes remain visible.
+- GPU work is verified by execution counters. No per-operator host
+  materialization/synchronization remains inside a supported dense region.
+- Tests cover values, dimensions, coordinates, attributes, bags, NULLs,
+  empty groups, defaults, masked exceptional values, grad/VJP/JVP, typed
+  refusals, re-registration and repeated updates. Existing tolerances stay.
+- Reports include cold/compile/setup and warm times, import/export costs,
+  memory, actual device work, fallback coverage and all five checked steps.
+
+An initial engineering target is at least a 10× reduction of the current
+123M warm-step time, followed by reducing the remaining gap to native
+array execution. This is a target, not a measured or promised speedup.
+Separate compatible Arrow API costs from optional resident-state costs.
+
+## Resume details
+
+Work directly in `/Users/cromwellian/xql/xarray-sql`, preserving existing
+uncommitted work. Read its `AGENTS.md`; when modifying the EinFold patch,
+read `/private/tmp/xql-einfold/AGENTS.md` and the relevant design sections.
+The user explicitly authorizes local implementation of this plan after
+06:00 Pacific. No publishing, merging or messages to maintainers are
+requested. Keep design changes and the reproducible local patch together.
+
+The checkout uses `.venv` Python 3.12.14, DataFusion Python 55.0.0,
+ddxdb 0.4.6 and custom EinFold `0.2.0+xqlmlx.3`, pinned to
+`39b25210d3b39027ef91584383fb114dfa8a4a48`. Native Rust DataFusion is 55.2.0.
+Use `PYTHONPATH=$PWD`; installed Python sources can differ. Do not run
+`uv sync`, which can replace the custom wheel. Metal access requires the
+appropriate execution permissions. Use `MLX_ENABLE_TF32=0` before import.
+
+EinFold sources/build cache: `/private/tmp/xql-einfold`; clean pinned tree:
+`/private/tmp/xql-einfold-clean`. Isolated Rust toolchain:
+`PATH=/private/tmp/xql-cargo/bin:$PATH`,
+`RUSTUP_HOME=/private/tmp/xql-rustup`,
+`CARGO_HOME=/private/tmp/xql-cargo`. Build EinFold from its `python`
+directory with maturin, locked dependencies and the checkout's Python.
+Retain all new native changes in the EinFold repository and verify that the pinned build
+applies cleanly to the pinned tree. Existing xarray-sql native builds must
+retain Flight support. No build processes remain paused.
+
+The upstream nanoGPT model and isolated PyTorch 2.15.0 environment are in
+`/private/tmp/xql-upstream-nanogpt` and `/private/tmp/xql-framework-venv`.
+The latter needs the checkout's `.venv/lib/python3.12/site-packages` on
+`PYTHONPATH` for JAX/NumPy. Baseline scripts are
+`benchmarks/apple/run_nanogpt_sizes.py` and
+`benchmarks/apple/compare_nanogpt_frameworks.py`. The last full local suite
+passed 637 tests (277 skipped, 102 integration deselected); native suites
+passed 60, with four upstream ignored tests and one unchanged plain-
+DataFusion ARM NaN-sign assumption excluded. Keep those qualifications.
diff --git a/accelerators/einfold_mlx/resident-training-plan.md b/accelerators/einfold_mlx/resident-training-plan.md
new file mode 100644
index 00000000..7ac1e8ec
--- /dev/null
+++ b/accelerators/einfold_mlx/resident-training-plan.md
@@ -0,0 +1,132 @@
+# Persistent GPU execution of SQL training
+
+The execution changes in this plan are implemented in
+[`TrainingSession`](../../xarray_sql/training.py) and the optional
+EinFold `.6` / ddx `.3` builds. The original plan below records their
+implementation order and acceptance criteria. Whole-step compilation,
+persistent state, shared primal loss, GPU dropout and explicit scalar
+validation are available in the Shakespeare trainer. Final validation and
+measurements are recorded in the resident execution report.
+
+[ddx](https://github.com/xqlsystems/ddx) differentiates SQL;
+[EinFold](https://github.com/xqlsystems/einfold) lowers numerical query
+plans to arrays; [MLX](https://github.com/ml-explore/mlx) executes the
+supported array arithmetic on Apple silicon. The goal is to retain this
+pipeline while keeping numerical training state in MLX between updates.
+
+## Current boundaries
+
+The [Shakespeare trainer](../../benchmarks/ml/03_shakespeare.py) calls a
+loss query, exports all gradients, registers them again, exports the
+clipping coefficient, then exports and registers weights and AdamW
+moments. Each `DenseExecutor._invoke` clears its transient tensor store
+and creates a new computed-value cache. These are deliberate properties
+of the current live-Dataset/Arrow snapshot API.
+
+The current [measurements](../../benchmarks/apple/shakespeare-training.md)
+are 2.024 s per warm SQL GPU update versus 0.02935 s in native PyTorch MPS
+for the 10.7M-parameter recipe. This plan does not attribute the entire
+difference to exports; Python plan interpretation, memory allocation,
+kernel dispatch and numerical operation selection also remain.
+
+Apple silicon uses [unified memory](https://ml-explore.github.io/mlx/build/html/usage/unified_memory.html).
+The costs to eliminate are conversion, packing, allocation and forced
+synchronization, rather than assuming every CPU/GPU boundary transfers
+data over PCIe.
+
+## Implementation order
+
+1. **Retain device tables within a complete training step.** Add an
+   opt-in internal execution path returning compact tables whose numerical
+   columns are MLX arrays. Execute gradients, global clipping and every
+   existing SQL AdamW statement in one invocation before exporting results.
+   Bind temporary gradient/state names to these actual device tables.
+   Schema providers supply planning metadata and must never substitute for
+   numerical evaluation or required validation. This first phase can still
+   export the final updated state, establishing correctness independently
+   of persistence.
+
+2. **Expose the primal value from ddx's backward program.** In the pinned
+   ddx release, `BackwardProgram.value` already identifies the loss output;
+   the Python binding exposes it. Read that output before transient tables
+   are released, together with the existing gradients. If the high-level
+   SQL preparation API does not expose the corresponding output handle,
+   add that bridge or execute the public backward program directly.
+   Do not assume adding an independent SELECT to a batch automatically
+   shares its forward computation. The loss and backward must use the same
+   saved forward values and dropout masks.
+
+3. **Add explicit persistent state ownership.** A new training-state
+   object owns weight and moment arrays, dtype/layout metadata, coordinate
+   domains, version and RNG state. Import the initial arrays once. Bind
+   the next update's SQL table names directly to the preceding update's
+   device outputs. Retain persistent state separately from transient
+   activations/cotangents; release temporary graphs after each step.
+   Publish all updated tables together only after checks succeed, retaining
+   the previous version on failure. Do not silently change the existing
+   API's handling of caller edits to NumPy buffers. Explicit replacement,
+   snapshot/export and close operations define the new ownership boundary.
+
+4. **Compile the complete numerical program.** Lower the prepared
+   forward/backward/optimizer DAG to stable array instructions once rather
+   than walking hundreds of DataFusion plans each update. Compile that
+   program with explicit weights, moments, batch, learning-rate/step scalars
+   and RNG inputs. Outputs include next state and the primal loss.
+   Cache by SQL/program identity, schema, coordinate/layout facts, shapes,
+   precision and UDF versions; invalidate when these change. Changing
+   scalar values must not trigger a compilation or capture stale weights.
+   Whole-program compilation can fuse pointwise work and share expressions;
+   it does not imply that every GEMM becomes one GPU kernel.
+
+5. **Keep masks, reductions and checks on the device.** Generate dropout
+   masks on MLX with checkpointed RNG state. Keep an explicit-mask mode for
+   exact comparisons against the current trainer/native oracle. Compute
+   the global gradient norm and all AdamW arithmetic on GPU. Prove static
+   coordinate uniqueness/domain properties from immutable metadata where
+   possible. Otherwise execute checks on actual device values and read a
+   scalar status before publishing an update. Preserve errors, NULL and
+   row-domain behavior; do not skip validation because providers have zero
+   rows. Unsupported plans decline before persistent state changes.
+
+6. **Make synchronization explicit.** Evaluate the step's device outputs
+   together and release references to previous computational histories.
+   Read the loss when logging rather than before backward. Export complete
+   arrays only for checkpoints, explicit inspection and correctness checks.
+   A required scalar validation result may still synchronize once per
+   update. Timed benchmarks must wait for device completion; an asynchronous
+   submission time is not an update time.
+
+MLX documents [whole-training-graph compilation and explicit state](https://ml-explore.github.io/mlx/build/html/usage/compile.html#compiling-training-graphs),
+including RNG state for dropout. Its examples use MLX differentiation;
+this executor would instead compile the array operations produced from
+ddx's SQL derivatives. The same compilation mechanism can be used without
+replacing the derivative source. MLX also documents how
+[scalar extraction forces evaluation](https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html).
+
+## Contract and acceptance criteria
+
+- SQL remains the source of the forward model, clipping and optimizer;
+  ddx remains the source of derivatives. Dense-only assumptions, precision
+  and ownership are explicit. Ordinary Arrow-returning SQL APIs continue
+  to produce independent snapshots.
+- Loss, raw gradients, weights, both moments and evaluation results match
+  the existing float64 CPU and native reference checks at the established
+  float32 tolerance. Exports for checks occur outside the performance timer.
+- Tests cover mask reuse, fresh parameters after every update, all-table
+  atomic publication, invalid row/domain checks, UDF/schema invalidation,
+  independent snapshots, explicit host replacement and checkpoint resume.
+- After initialization, the resident step performs no parameter-,
+  moment- or gradient-sized host exports/imports and creates no Arrow
+  gradient/update tables. Batch data and optional scalar control/status
+  are the remaining routine boundaries. Count both bytes and sync events.
+- Report staged timings: current baseline; combined device step; reused
+  primal loss; persistent state; compiled program. Measure compilation
+  separately. Compare five checked steps and at least 100 steady-state
+  updates, including synchronized duration and peak memory. Confirm bounded
+  memory and unchanged SQL/gradient semantics before claiming speedups.
+
+If a substantial gap remains after these boundaries are removed, profile
+the resulting compiled execution. Candidate work includes backward
+contraction scheduling, common expressions, activation liveness and fused
+normalization/attention kernels selected from proven plan patterns.
+Native MPS parity cannot be predicted from the current measurements.
diff --git a/benchmarks/ml/03_shakespeare.py b/benchmarks/ml/03_shakespeare.py
new file mode 100644
index 00000000..946c2a9a
--- /dev/null
+++ b/benchmarks/ml/03_shakespeare.py
@@ -0,0 +1,714 @@
+"""Train nanoGPT's Shakespeare character model with SQL and ddx gradients.
+
+nanoGPT (https://github.com/karpathy/nanoGPT) supplies the training recipe
+and prepared character data. The model, loss, gradient clipping and AdamW
+updates execute SQL; EinFold (https://github.com/xqlsystems/einfold) lowers
+the plans to compact CPU arrays or MLX (https://github.com/ml-explore/mlx)
+GPU arrays. Python samples batches and dropout masks and writes checkpoints.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import math
+import os
+import resource
+import subprocess
+import time
+from pathlib import Path
+
+import ddxdb
+import numpy as np
+import xarray as xr
+from _hosts import DenseCPU, DenseGPU
+from _shakespeare_torch import _TorchReference
+
+import xarray_sql as xql
+from xarray_sql.math import register_erf
+from xarray_sql.training import TrainingSession
+
+_SPEC = importlib.util.spec_from_file_location(
+    "_shakespeare_sql_model", Path(__file__).with_name("02_nanogpt.py")
+)
+_MODEL = importlib.util.module_from_spec(_SPEC)
+_SPEC.loader.exec_module(_MODEL)
+
+# The maximum cancels out of cross entropy. Spell its reduction as a join
+# to avoid repeatedly traversing a shared model graph during window lowering.
+# Holding this shared shift constant removes only a cancelling derivative.
+_MODEL.LOSS = """
+predicted AS (SELECT seq,pos,tok,val FROM logits),
+maximum AS (SELECT seq,pos,ddx_stop_gradient(MAX(val)) AS m FROM predicted GROUP BY seq,pos),
+log_sum_exp AS (
+  SELECT p.seq,p.pos,MAX(m.m)+ln(SUM(exp(p.val-m.m))) AS lse
+  FROM predicted p JOIN maximum m ON p.seq=m.seq AND p.pos=m.pos
+  GROUP BY p.seq,p.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 _parser():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument(
+        "--backend", choices=["dense_cpu", "dense_gpu"], default="dense_gpu"
+    )
+    p.add_argument(
+        "--data-dir",
+        type=Path,
+        default=Path.home() / "nanoGPT/data/shakespeare_char",
+    )
+    p.add_argument("--out-dir", type=Path, default=Path("out-shakespeare-sql"))
+    p.add_argument(
+        "--resume",
+        type=Path,
+        help="resume an NPZ checkpoint including moments and RNGs",
+    )
+    p.add_argument(
+        "--steps",
+        type=int,
+        help="cap this run's updates without shortening the LR schedule",
+    )
+    p.add_argument(
+        "--verify-nanogpt",
+        type=Path,
+        help="check every gradient, update and moment against this nanoGPT checkout",
+    )
+    for name, default in {
+        "batch_size": 12,
+        "block_size": 256,
+        "n_layer": 6,
+        "n_head": 6,
+        "n_embd": 384,
+        "gradient_accumulation_steps": 1,
+        "max_iters": 2000,
+        "lr_decay_iters": 2000,
+        "warmup_iters": 100,
+        "eval_interval": 200,
+        "eval_iters": 20,
+        "log_interval": 10,
+        "seed": 1337,
+    }.items():
+        p.add_argument(
+            "--" + name,
+            "--" + name.replace("_", "-"),
+            type=int,
+            default=default,
+        )
+    for name, default in {
+        "dropout": 0.2,
+        "learning_rate": 1e-3,
+        "min_lr": 1e-4,
+        "beta1": 0.9,
+        "beta2": 0.99,
+        "weight_decay": 0.1,
+        "grad_clip": 1.0,
+    }.items():
+        p.add_argument(
+            "--" + name,
+            "--" + name.replace("_", "-"),
+            type=float,
+            default=default,
+        )
+    p.add_argument("--cache-gib", type=float, default=16)
+    p.add_argument(
+        "--resident",
+        action="store_true",
+        help="keep SQL training state as owned device arrays",
+    )
+    p.add_argument(
+        "--compile-training",
+        action="store_true",
+        help="compile the complete SQL-derived device step",
+    )
+    p.add_argument(
+        "--gpu-dropout",
+        action="store_true",
+        help="generate dropout masks inside the device training step",
+    )
+    p.add_argument(
+        "--profile-memory",
+        action="store_true",
+        help="sample current RSS outside training timers at log intervals",
+    )
+    return p
+
+
+def _learning_rate(it, args):
+    if it < args.warmup_iters:
+        return args.learning_rate * (it + 1) / (args.warmup_iters + 1)
+    if it > args.lr_decay_iters:
+        return args.min_lr
+    ratio = (it - args.warmup_iters) / (args.lr_decay_iters - args.warmup_iters)
+    return args.min_lr + 0.5 * (1 + math.cos(math.pi * ratio)) * (
+        args.learning_rate - args.min_lr
+    )
+
+
+def _data(args):
+    # Match prepare.py's uint16 files and sorted character vocabulary without
+    # executing the pickle in meta.pkl. Check that the files encode input.txt.
+    text = (args.data_dir / "input.txt").read_text()
+    chars = sorted(set(text))
+    stoi = {c: i for i, c in enumerate(chars)}
+    encoded = np.fromiter((stoi[c] for c in text), dtype=np.uint16)
+    split = int(len(encoded) * 0.9)
+    data = {}
+    for name, expected in (
+        ("train", encoded[:split]),
+        ("val", encoded[split:]),
+    ):
+        values = np.memmap(
+            args.data_dir / f"{name}.bin", dtype=np.uint16, mode="r"
+        )
+        if not np.array_equal(values, expected):
+            raise ValueError(
+                f"{name}.bin differs from nanoGPT's character preparation of input.txt"
+            )
+        if len(values) <= args.block_size:
+            raise ValueError(f"{name} split is shorter than the context")
+        data[name] = values
+    return chars, data
+
+
+def _batch(values, args, rng):
+    starts = rng.integers(len(values) - args.block_size, size=args.batch_size)
+    tok = np.stack([values[i : i + args.block_size] for i in starts]).astype(
+        np.int64
+    )
+    nxt = np.stack(
+        [values[i + 1 : i + args.block_size + 1] for i in starts]
+    ).astype(np.int64)
+    return xr.Dataset(
+        {"tok": (("seq", "pos"), tok), "next": (("seq", "pos"), nxt)},
+        coords={
+            "seq": np.arange(args.batch_size),
+            "pos": np.arange(args.block_size),
+        },
+    )
+
+
+def _masks(args, rng):
+    def mask(dims, shape):
+        scale = (rng.random(shape) >= args.dropout).astype(np.float64) / (
+            1 - args.dropout
+        )
+        return xr.Dataset(
+            {"scale": (dims, scale)},
+            coords={d: np.arange(n) for d, n in zip(dims, shape)},
+        )
+
+    shape = (args.batch_size, args.block_size, args.n_embd)
+    result = {"embed_drop": mask(("seq", "pos", "e"), shape)}
+    for layer in range(args.n_layer):
+        result[f"l{layer}_adrop"] = mask(
+            ("seq", "h", "i", "j"),
+            (args.batch_size, args.n_head, args.block_size, args.block_size),
+        )
+        for suffix in ("rdrop", "mdrop"):
+            result[f"l{layer}_{suffix}"] = mask(("seq", "pos", "e"), shape)
+    return result
+
+
+def _clip_sql(args):
+    if args.grad_clip == 0:
+        return "SELECT 1.0 AS c"
+    every_grad = " UNION ALL ".join(
+        f"SELECT g FROM {t}_grad" for t in _MODEL.PARAMS
+    )
+    return f"""WITH norm AS (SELECT sqrt(SUM(g*g)) AS n FROM ({every_grad})),
+    coefficient AS (SELECT {args.grad_clip}/(n+1e-6) AS c FROM norm)
+    SELECT CASE WHEN c < 1.0 THEN c ELSE 1.0 END AS c FROM coefficient"""
+
+
+def _update_sql(name, args):
+    dims = _MODEL.PARAMS[name]
+    keys = ", ".join(f"p.{d}" for d in dims)
+    same = lambda alias: " AND ".join(f"p.{d}={alias}.{d}" for d in dims)
+    decay = 0.0 if name.startswith("ln_") else args.weight_decay
+    return f"""WITH moments AS (
+    SELECT {keys}, p.w,
+      {args.beta1}*s.m + (1.0-{args.beta1})*g.g*c.c AS m,
+      {args.beta2}*s.v + (1.0-{args.beta2})*power(g.g*c.c,2) AS v
+    FROM {name} p JOIN {name}_grad g ON {same("g")}
+    JOIN {name}_state s ON {same("s")} CROSS JOIN clip c)
+    SELECT {keys},
+      p.w*(1.0-t.lr*{decay}) - t.lr*(p.m/t.b1)/(sqrt(p.v/t.b2)+1e-8) AS w,
+      p.m, p.v FROM moments p CROSS JOIN train_step t ORDER BY {keys}"""
+
+
+def _save(
+    path,
+    params,
+    moments,
+    iteration,
+    args,
+    chars,
+    data_rng,
+    mask_rng,
+    best,
+    skip_eval=False,
+):
+    metadata = {
+        "iteration": iteration,
+        "config": vars(args)
+        | {
+            "bias": False,
+            "dtype": "float32" if args.backend == "dense_gpu" else "float64",
+        },
+        "chars": chars,
+        "data_rng": data_rng.bit_generator.state,
+        "mask_rng": mask_rng.bit_generator.state,
+        "best_val_loss": best,
+        "skip_eval": skip_eval,
+    }
+    arrays = {
+        f"{name}_{col}": value
+        for name, w in params.items()
+        for col, value in zip(("w", "m", "v"), (w, *moments[name]))
+    }
+    temporary = path.with_suffix(".tmp")
+    with temporary.open("wb") as handle:
+        np.savez(
+            handle,
+            metadata=np.array(json.dumps(metadata, default=str)),
+            **arrays,
+        )
+    temporary.replace(path)
+
+
+def main():
+    parser = _parser()
+    args = parser.parse_args()
+    if args.gpu_dropout:
+        args.resident = True
+        if args.backend != "dense_gpu":
+            parser.error("--gpu-dropout requires dense_gpu")
+    if args.compile_training:
+        args.resident = True
+        if args.backend != "dense_gpu":
+            parser.error("--compile-training requires dense_gpu")
+    if (
+        min(
+            args.batch_size,
+            args.block_size,
+            args.n_layer,
+            args.n_head,
+            args.n_embd,
+            args.eval_interval,
+            args.eval_iters,
+            args.log_interval,
+        )
+        < 1
+        or args.n_embd % args.n_head
+    ):
+        parser.error(
+            "model sizes and evaluation/logging intervals must be positive, with n_embd divisible by n_head"
+        )
+    if args.gradient_accumulation_steps != 1:
+        parser.error(
+            "this recipe currently supports gradient_accumulation_steps=1"
+        )
+    if (
+        not 0 <= args.beta1 < 1
+        or not 0 <= args.beta2 < 1
+        or not all(
+            math.isfinite(v) and v >= 0
+            for v in (
+                args.learning_rate,
+                args.min_lr,
+                args.weight_decay,
+                args.grad_clip,
+            )
+        )
+        or not math.isfinite(args.cache_gib)
+        or args.cache_gib <= 0
+    ):
+        parser.error(
+            "require betas in [0,1), finite nonnegative optimizer values and a positive cache size"
+        )
+    if (
+        not 0 <= args.dropout < 1
+        or not 0 <= args.warmup_iters < args.lr_decay_iters
+        or args.max_iters < 0
+        or (args.steps is not None and args.steps < 1)
+    ):
+        parser.error(
+            "require 0 <= dropout < 1, 0 <= warmup_iters < lr_decay_iters and positive steps"
+        )
+    if args.backend == "dense_gpu" and os.environ.get("MLX_ENABLE_TF32") != "0":
+        parser.error(
+            "GPU execution requires MLX_ENABLE_TF32=0 before Python starts"
+        )
+    if "erf" not in ddxdb.supported_functions():
+        parser.error(
+            "exact GELU needs the local ddx math build: python accelerators/ddx_math/build.py"
+        )
+    _MODEL.LAYERS, _MODEL.EMBED, _MODEL.HEADS = (
+        args.n_layer,
+        args.n_embd,
+        args.n_head,
+    )
+    _MODEL.HEAD, _MODEL.HIDDEN, _MODEL.CONTEXT = (
+        args.n_embd // args.n_head,
+        4 * args.n_embd,
+        args.block_size,
+    )
+    chars, data = _data(args)
+    data_rng, mask_rng = (
+        np.random.default_rng(args.seed),
+        np.random.default_rng(args.seed + 1),
+    )
+    params = _MODEL.init_params(len(chars), np.random.default_rng(args.seed))
+    moments = {
+        name: (np.zeros_like(w), np.zeros_like(w)) for name, w in params.items()
+    }
+    iteration, best, skip_eval = 0, math.inf, False
+    if args.resume:
+        with np.load(args.resume, allow_pickle=False) as saved:
+            meta = json.loads(str(saved["metadata"]))
+            ignored = {
+                "steps",
+                "resume",
+                "out_dir",
+                "data_dir",
+                "verify_nanogpt",
+                "cache_gib",
+                "max_iters",
+                "profile_memory",
+                "resident",
+                "compile_training",
+            }
+            for key, value in vars(args).items():
+                if (
+                    key not in ignored
+                    and meta["config"].get(key, False) != value
+                ):
+                    parser.error(f"checkpoint configuration differs for {key}")
+            if meta["chars"] != chars:
+                parser.error("checkpoint vocabulary differs from input.txt")
+            for name, w in params.items():
+                for col in ("w", "m", "v"):
+                    if saved[f"{name}_{col}"].shape != w.shape:
+                        parser.error(
+                            f"checkpoint shape differs for {name}_{col}"
+                        )
+                params[name] = saved[f"{name}_w"].copy()
+                moments[name] = tuple(
+                    saved[f"{name}_{c}"].copy() for c in ("m", "v")
+                )
+            iteration, best, skip_eval = (
+                meta["iteration"],
+                meta["best_val_loss"],
+                meta["skip_eval"],
+            )
+            data_rng.bit_generator.state = meta["data_rng"]
+            mask_rng.bit_generator.state = meta["mask_rng"]
+    if args.verify_nanogpt and args.resume:
+        parser.error(
+            "reference verification starts from scratch; resume preserves SQL optimizer state"
+        )
+    os.environ["XQL_DENSE_CACHE_BYTES"] = str(int(args.cache_gib * (1 << 30)))
+    host = DenseGPU() if args.backend == "dense_gpu" else DenseCPU()
+    register_erf(host.ctx)
+    for name, w in params.items():
+        host.load(name, _MODEL.as_table(name, w))
+        host.load(
+            name + "_state",
+            _MODEL.as_table(name, w, *moments[name]).drop_vars("w"),
+        )
+    # Schema registration must not advance the checkpointed RNG streams.
+    host.load("batch", _batch(data["train"], args, np.random.default_rng(0)))
+    for name, value in _masks(args, np.random.default_rng(0)).items():
+        host.load(name, value)
+    host.execute(
+        _MODEL.model(exact_gelu=True, dropout=True, grouped_softmax=True)
+    )
+    # Register the second view directly: both graphs persist across evaluation.
+    eval_model = _MODEL.model(
+        exact_gelu=True, view="eval_logits", grouped_softmax=True
+    )
+    host.executor.register_view("eval_logits", eval_model.split(" AS\n", 1)[1])
+    resident = None
+    resident_step = None
+    if args.resident:
+        resident = TrainingSession(
+            host.ctx,
+            precision=host.dense_precision,
+            max_cached_bytes=int(args.cache_gib * (1 << 30)),
+            inline_ctes=host.executor.inline_ctes,
+        )
+        for name, query in host.executor._views.items():
+            resident.executor.register_view(name, query)
+        for name in host._tables:
+            resident.bind(name)
+        if args.compile_training:
+            resident.index_domains(
+                "batch", {"tok": ("wte", "tok"), "next": ("wte", "tok")}
+            )
+            resident.bind("batch")
+        if args.gpu_dropout:
+            resident.dropout(
+                [name for name in host._tables if name.endswith("drop")],
+                probability=args.dropout,
+                seed=args.seed + 1,
+                step=iteration,
+            )
+        updates_sql = {
+            name + "_next": _update_sql(name, args) for name in params
+        }
+        assignments = {}
+        for name, dims in _MODEL.PARAMS.items():
+            assignments[name] = (name + "_next", [*dims, "w"])
+            assignments[name + "_state"] = (name + "_next", [*dims, "m", "v"])
+        resident_step = resident.prepare_step(
+            gradients=dict(
+                zip(
+                    (name + "_grad" for name in params),
+                    _MODEL.gradient_statements(),
+                )
+            ),
+            stages={"clip": _clip_sql(args)},
+            updates=updates_sql,
+            assignments=assignments,
+            compile=args.compile_training,
+        )
+
+    def load(name, value):
+        if resident is None:
+            host.load(name, value)
+        else:
+            resident.bind(name, value)
+
+    def download():
+        if resident is None:
+            return
+        tables = resident.export(*params, *(name + "_state" for name in params))
+        for name in params:
+            params[name] = _MODEL.as_array(name, tables[name])
+            state = xql.to_dataset(
+                tables[name + "_state"], dims=list(_MODEL.PARAMS[name])
+            )
+            moments[name] = (state.m.values, state.v.values)
+
+    train_loss = host.executor.prepare(_MODEL.loss_statement())
+    eval_loss = host.executor.prepare(
+        _MODEL.loss_statement().replace("FROM logits", "FROM eval_logits")
+    )
+    gradients = host.executor.prepare_all(_MODEL.gradient_statements())
+    clip = host.executor.prepare(_clip_sql(args))
+    updates = {
+        name: host.executor.prepare(_update_sql(name, args)) for name in params
+    }
+    reference = (
+        _TorchReference(args, params, len(chars))
+        if args.verify_nanogpt
+        else None
+    )
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    print(
+        json.dumps(
+            {
+                "backend": args.backend,
+                "parameters": sum(w.size for w in params.values()),
+                "vocabulary": len(chars),
+                "config": vars(args),
+            },
+            default=str,
+        ),
+        flush=True,
+    )
+    # train.py tests iter_num > max_iters after each update, so iteration
+    # max_iters is included. --steps provides an exact short-run update cap.
+    end = args.max_iters + 1
+    if args.steps is not None:
+        end = min(end, iteration + args.steps)
+    with (args.out_dir / "metrics.jsonl").open(
+        "a" if args.resume else "w"
+    ) as log:
+        while iteration < end:
+            if iteration % args.eval_interval == 0 and not skip_eval:
+                losses = {}
+                for split in ("train", "val"):
+                    measured = []
+                    for _ in range(args.eval_iters):
+                        batch = _batch(data[split], args, data_rng)
+                        load("batch", batch)
+                        value = float(
+                            (
+                                eval_loss()
+                                if resident is None
+                                else resident.query(
+                                    _MODEL.loss_statement().replace(
+                                        "FROM logits", "FROM eval_logits"
+                                    )
+                                )
+                            )["l"][0].as_py()
+                        )
+                        if reference:
+                            reference.evaluate(batch, value)
+                        measured.append(value)
+                    losses[split] = float(np.mean(measured))
+                event = dict(kind="eval", iteration=iteration, **losses)
+                print(json.dumps(event), flush=True)
+                log.write(json.dumps(event) + "\n")
+                log.flush()
+                if losses["val"] < best:
+                    best = losses["val"]
+                    if iteration > 0:
+                        download()
+                        _save(
+                            args.out_dir / "best.npz",
+                            params,
+                            moments,
+                            iteration,
+                            args,
+                            chars,
+                            data_rng,
+                            mask_rng,
+                            best,
+                            skip_eval=True,
+                        )
+            skip_eval = False
+            batch = _batch(data["train"], args, data_rng)
+            masks = {} if args.gpu_dropout else _masks(args, mask_rng)
+            load("batch", batch)
+            for name, value in masks.items():
+                load(name, value)
+            lr = _learning_rate(iteration, args)
+            load(
+                "train_step",
+                xr.Dataset(
+                    {
+                        "lr": xr.DataArray(lr),
+                        "b1": xr.DataArray(1 - args.beta1 ** (iteration + 1)),
+                        "b2": xr.DataArray(1 - args.beta2 ** (iteration + 1)),
+                    }
+                ),
+            )
+            started = time.perf_counter()
+            if resident is None:
+                loss = float(train_loss()["l"][0].as_py())
+                grad_rows = gradients()
+                grads = {
+                    name: xql.to_dataset(rows, dims=list(_MODEL.PARAMS[name]))
+                    for name, rows in zip(params, grad_rows)
+                }
+                for name, ds in grads.items():
+                    host.load(name + "_grad", ds)
+                host.load("clip", xql.to_dataset(clip(), dims=[]))
+                # Compute every update before replacing any current weight.
+                next_tables = {
+                    name: xql.to_dataset(run(), dims=list(_MODEL.PARAMS[name]))
+                    for name, run in updates.items()
+                }
+                for name, ds in next_tables.items():
+                    params[name] = ds.w.values
+                    moments[name] = (ds.m.values, ds.v.values)
+                    host.load(name, ds[["w"]])
+                    host.load(name + "_state", ds[["m", "v"]])
+            else:
+                result = resident_step()
+                loss = result.loss_value()
+            seconds = time.perf_counter() - started
+            if reference:
+                if resident is not None:
+                    if args.gpu_dropout:
+                        mask_tables = result.export(
+                            *(
+                                name
+                                for name in host._tables
+                                if name.endswith("drop")
+                            )
+                        )
+                        masks = {
+                            name: xql.to_dataset(
+                                rows,
+                                dims=[
+                                    f for f in rows.column_names if f != "scale"
+                                ],
+                            )
+                            for name, rows in mask_tables.items()
+                        }
+                    grad_tables = result.export(
+                        *(name + "_grad" for name in params)
+                    )
+                    grads = {
+                        name: xql.to_dataset(
+                            grad_tables[name + "_grad"],
+                            dims=list(_MODEL.PARAMS[name]),
+                        )
+                        for name in params
+                    }
+                    download()
+                reference.step(batch, masks, lr, loss, grads, params, moments)
+            event = {
+                "kind": "train",
+                "iteration": iteration,
+                "loss": loss,
+                "lr": lr,
+                "seconds": seconds,
+                "tokens_per_second": args.batch_size
+                * args.block_size
+                / seconds,
+                "verified": reference is not None,
+            }
+            if resident is not None:
+                event["device_step"] = result.measurements
+            if args.profile_memory and (
+                iteration % args.log_interval == 0 or iteration + 1 == end
+            ):
+                event["rss_bytes"] = (
+                    int(
+                        subprocess.check_output(
+                            ["ps", "-p", str(os.getpid()), "-o", "rss="],
+                            text=True,
+                        ).strip()
+                    )
+                    * 1024
+                )
+                event["peak_rss_bytes"] = resource.getrusage(
+                    resource.RUSAGE_SELF
+                ).ru_maxrss
+                if (
+                    resident is not None
+                    and resident.executor._runtime is not None
+                ):
+                    event["mlx_memory"] = resident.executor._runtime.memory()
+            log.write(json.dumps(event) + "\n")
+            log.flush()
+            if iteration % args.log_interval == 0 or iteration + 1 == end:
+                print(json.dumps(event), flush=True)
+            iteration += 1
+    download()
+    _save(
+        args.out_dir / "last.npz",
+        params,
+        moments,
+        iteration,
+        args,
+        chars,
+        data_rng,
+        mask_rng,
+        best,
+    )
+    print(
+        json.dumps(
+            {
+                "checkpoint": str(args.out_dir / "last.npz"),
+                "next_iteration": iteration,
+                "measurements": host.measurements()
+                if resident is None
+                else resident.statistics(),
+            }
+        ),
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md
index ee632f96..651dd8eb 100644
--- a/benchmarks/ml/README.md
+++ b/benchmarks/ml/README.md
@@ -1,5 +1,13 @@
 # Machine learning in SQL
 
+The optional `dense_cpu` and `dense_gpu` hosts retain compact arrays across
+SQL and gradient stages using the local
+[EinFold integration](../../accelerators/einfold_mlx/README.md).
+`dense_gpu` explicitly permits float32 arithmetic and requires
+`MLX_ENABLE_TF32=0`. The hosts use identical model SQL, ddx derivatives and
+existing reference checks. See the
+[dense execution measurements](../apple/dense-execution.md).
+
 These scripts write 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:
@@ -11,7 +19,7 @@ same SQL on several engines. Two related projects shape them:
   are its [target demos](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md),
   and the demos run on DataFusion with einfold as one of their engines.
 
-Each script checks its results against the same model in
+The first two scripts check their results against the same model in
 [JAX](https://github.com/jax-ml/jax), Google's library for differentiable
 array programs.
 
@@ -19,6 +27,7 @@ array programs.
 |---|---|---|---|
 | 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, against [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan, which einfold finds from the same SQL; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention: the [Nadaraya–Watson estimator](https://en.wikipedia.org/wiki/Kernel_regression#Nadaraya%E2%80%93Watson_kernel_regression), an average of nearby points weighted by a Gaussian kernel, which einfold must fuse with the same rules |
 | 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_shakespeare.py`](03_shakespeare.py) | nanoGPT's 6-layer, 384-channel, 6-head Shakespeare character model with context 256 and batch 12 | training from scratch with exact GELU, dropout, SQL AdamW, train/validation evaluation and resumable checkpoints; optional checks of every gradient, weight and optimizer moment against the original PyTorch model |
 
 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
@@ -29,7 +38,7 @@ these models fast is tracked in
 [#285](https://github.com/xqlsystems/xarray-sql/issues/285), which lists
 places to start.
 
-Every result is checked before it is timed into the table. Each script
+Every result in the original demos is checked before it is timed into the table. Each demo
 computes the same model in JAX from the same inputs and parameters, and a
 mismatch fails the run.
 
@@ -40,12 +49,99 @@ uv run benchmarks/ml/01_attention.py
 uv run benchmarks/ml/02_nanogpt.py --steps 5
 ```
 
-Each script carries [inline `uv` metadata](https://docs.astral.sh/uv/guides/scripts/)
+### Shakespeare training matching nanoGPT's train.py
+
+`03_shakespeare.py` is a training loop with checkpoints and held-out
+validation. Its defaults match `train_shakespeare_char.py` with these
+overrides: batch 12, context 256, layers/heads/channels 6/6/384, gradient
+accumulation 1, `max_iters=2000`, `lr_decay_iters=2000`, warmup 100,
+evaluation every 200 iterations over 20 batches per split. Inherited
+settings also match: dropout 0.2, exact GELU, no biases, tied embeddings,
+AdamW with learning rate 1e-3 to 1e-4, betas 0.9/0.99, epsilon 1e-8,
+weight decay 0.1 excluding layer norm gains, and global gradient clipping
+at 1. The learning rate uses `lr*(iteration+1)/(warmup_iters+1)` during
+warmup, as in the local nanoGPT source. There are **10,745,088 parameters**
+with the 65-character vocabulary; nanoGPT prints 10.65M when excluding
+position embeddings.
+
+Prepare the dataset once in a nanoGPT checkout with
+`python data/shakespeare_char/prepare.py`. The SQL script reads its
+`input.txt`, `train.bin` and `val.bin`, verifies the same sorted vocabulary
+and 90/10 split, and samples shifted next-character targets without
+crossing split boundaries. It initializes weights once per invocation
+and keeps updating them across iterations.
+
+Use this checkout's existing virtual environment with the optional
+[EinFold build](../../accelerators/einfold_mlx/README.md) and
+[ddx erf build](../../accelerators/ddx_math/README.md). Avoid `uv run`
+for this script, since environment synchronization can replace those local
+wheels. From the xarray-sql directory:
+
+```sh
+MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/03_shakespeare.py \
+  --backend dense_gpu \
+  --data-dir /Users/cromwellian/nanoGPT/data/shakespeare_char \
+  --out-dir out-shakespeare-sql
+```
+
+Use `--backend dense_cpu --cache-gib 32` for the measured compact float64 CPU path. GPU uses MLX
+float32. Both differentiate the original SQL with ddx and execute the
+model, global gradient clipping and AdamW arithmetic as SQL. They cache
+query programs across updates. AdamW moments live in separate SQL tables
+so the model's differentiated parameter column remains `w`. The maximum
+used for softmax stabilization is spelled with `GROUP BY` and joins to
+avoid repeated window-lowering traversal through the shared model graph.
+Its maximum is held constant with `ddx_stop_gradient`: subtracting any
+shared constant leaves softmax and cross entropy unchanged, so this
+removes a cancelling derivative through the maximum without changing the
+model's loss or parameter gradients.
+
+All size and recipe options accept nanoGPT's underscore spelling and
+hyphens, for example `--batch_size=12` and `--eval-iters=20`. The loop
+matches `train.py`'s termination condition: `max_iters=2000` includes
+iterations 0 through 2000, or 2001 optimizer updates. Add `--steps 5` to
+run exactly five updates while retaining the 2000-iteration LR schedule.
+
+Evaluation runs before the update at iteration 0, then every 200
+iterations, with dropout disabled on 20 train and 20 validation batches.
+`metrics.jsonl` records losses, learning rates and synchronized SQL
+training times. Times include forward loss, backward gradients, clipping,
+AdamW and result registration; sampling masks/batches, evaluation,
+checkpoint I/O and optional reference checks are outside the timer.
+
+`last.npz` contains all weights, optimizer moments, next iteration, both
+RNG states, vocabulary and configuration. `best.npz` saves an improved
+validation result after iteration 0. Resume with the same configuration:
+
+```sh
+MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/03_shakespeare.py \
+  --resume out-shakespeare-sql/last.npz
+```
+
+For an independent correctness check, install PyTorch in the execution
+environment and add `--steps 5 --verify-nanogpt /Users/cromwellian/nanoGPT`.
+This imports the original `model.py`, retains exact GELU and PyTorch
+AdamW, and compares evaluation/training loss, every raw gradient, every
+updated weight and both moments on every step. The oracle uses explicit
+attention and the same externally sampled Bernoulli dropout masks;
+reference checks do not replace SQL computation. GPU comparisons use
+the existing 3e-5 absolute/relative tolerance and CPU comparisons 1e-8.
+
+Random initialization and sampling use NumPy's RNG rather than PyTorch's,
+so the same integer seed alone does not reproduce a separate stock
+`train.py` trajectory. The oracle loads identical initial arrays and
+masks to test the arithmetic. Checkpoints use this script's safe NPZ
+format, rather than nanoGPT's PyTorch checkpoint format. This script
+currently supports gradient accumulation 1 and training without biases.
+The [measured comparison](../apple/shakespeare-training.md) records the
+verified five-step CPU/GPU runs and validation results.
+
+The first two scripts carry [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 each:
 
 ```shell
---hosts datafusion duckdb   # the engines to run on (default: all)
+--hosts datafusion duckdb   # default: datafusion einfold duckdb
 --sizes 64 256              # problem sizes: tokens, or batch size
 --reps 3                    # timed repetitions; the median is reported
 --csv results.csv           # also write the results table
@@ -105,6 +201,58 @@ in its plan.
 | `einfold` | `xarray_sql.Context`, with einfold as its query planner | ddx, whose steps einfold plans too |
 | `duckdb` | DuckDB, through `xql.register` | ddx, with the `duckdb` extra of ddxdb, which pins the DuckDB it supports |
 
+The optional Apple silicon comparison adds explicit hosts (the three
+defaults above stay the same):
+
+| Host | Execution |
+|---|---|
+| `datafusion_prepared` | Retain forward query plans, ddx programs, model views and parameter buffers |
+| `cpu_optimized` | Prepared path with EinFold and the float64 CPU matrix backend, using its normal thresholds |
+| `cpu_forced` | Same, with the backend threshold set to zero for a direct small-kernel comparison |
+| `gpu_optimized` | Prepared path with [MLX](https://github.com/ml-explore/mlx), Apple's array framework, and measured automatic GPU dispatch thresholds |
+| `gpu_forced` | Same, forcing eligible dense operations to GPU to measure dispatch overhead |
+| `shared_einfold` | Shared forward CTEs and reusable ddx stages, with EinFold's CPU kernels |
+| `shared_cpu` | Shared execution with the float64 CPU matrix backend |
+| `shared_gpu` | Shared execution with explicitly enabled float32 MLX kernels |
+
+These hosts need the [optional native bridge](../../accelerators/einfold_mlx/README.md).
+GPU hosts explicitly permit float32 arithmetic and check JAX with
+`rtol=atol=3e-5`; float64 hosts retain `1e-9`. Completed CPU/GPU kernel
+counts are recorded, so an automatic-dispatch run that stays on CPU is
+visible. Dataset loading and model-view creation times are also recorded
+cumulatively on training rows. The loss/training times keep their original
+scope and exclude JAX reference work and those loading/view times.
+
+```sh
+MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/ml/02_nanogpt.py \
+  --steps 5 --hosts datafusion einfold duckdb datafusion_prepared \
+  cpu_optimized cpu_forced gpu_optimized gpu_forced --csv comparison.csv
+```
+
+The [M5 Ultra five-step comparison](../apple/nanogpt-5-steps.md) reports
+cold and warm costs separately, checks every loss and weight update, and
+includes a forced GPU run with actual kernel counts.
+
+The [shared-execution report](../apple/nanogpt-shared.md) adds larger models.
+`--preset demo` uses one layer, 128 channels, four heads and context 64;
+`--preset shakespeare` uses six layers, 384 channels, six heads and context
+256, following nanoGPT's character-model dimensions. `--preset gpt2` uses
+12 layers, 768 channels, 12 heads, context 1024 and vocabulary 50257.
+Explicit dimension flags override presets. Inputs remain Shakespeare
+character IDs: padding the vocabulary does not add GPT-2 tokenization,
+pretrained weights, biases or dropout. `--forward-only` checks losses
+without running gradients or updates.
+
+```sh
+MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python benchmarks/apple/run_nanogpt_sizes.py \
+  --preset shakespeare --steps 5 --batch 1 --learning-rate 0.01 \
+  --hosts shared_cpu shared_gpu --prefix local-shakespeare
+```
+
+This runner starts hosts serially, checkpoints checked CSV rows and records
+commands, exit status and sampled resident memory. A failed or bounded run
+is not reported as a completed training benchmark.
+
 A host that cannot differentiate still runs every forward query, and its
 gradient rows say `no ddx backend`. A gradient that ddx, or the engine
 it plans with, refuses is reported as `refused:` with the error, and a
diff --git a/benchmarks/ml/_shakespeare_torch.py b/benchmarks/ml/_shakespeare_torch.py
new file mode 100644
index 00000000..c9c742e8
--- /dev/null
+++ b/benchmarks/ml/_shakespeare_torch.py
@@ -0,0 +1,214 @@
+"""Optional numerical oracle using the original nanoGPT model and AdamW.
+
+nanoGPT (https://github.com/karpathy/nanoGPT) runs the same batches,
+initial weights and externally supplied Bernoulli dropout masks. Only
+dropout sampling and the fused attention selection are adapted, allowing
+each SQL gradient and optimizer state to be checked independently.
+"""
+
+import importlib.util
+import sys
+
+import numpy as np
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+class _TorchReference:
+    def __init__(self, args, params, vocab):
+        if torch is None:
+            raise ImportError(
+                "--verify-nanogpt requires PyTorch in this Python environment"
+            )
+        spec = importlib.util.spec_from_file_location(
+            "_shakespeare_native_model", args.verify_nanogpt / "model.py"
+        )
+        native = importlib.util.module_from_spec(spec)
+        sys.modules[spec.name] = native
+        spec.loader.exec_module(native)
+        torch.set_num_threads(8)
+        self.args = args
+        self.dtype = (
+            torch.float32 if args.backend == "dense_gpu" else torch.float64
+        )
+        self.tolerance = 3e-5 if self.dtype == torch.float32 else 1e-8
+        self.model = native.GPT(
+            native.GPTConfig(
+                block_size=args.block_size,
+                vocab_size=vocab,
+                n_layer=args.n_layer,
+                n_head=args.n_head,
+                n_embd=args.n_embd,
+                dropout=args.dropout,
+                bias=False,
+            )
+        ).to(dtype=self.dtype)
+        self.masks = {}
+        owner = self
+
+        class _Mask(torch.nn.Module):
+            def __init__(self, name):
+                super().__init__()
+                self.name = name
+
+            def forward(self, x):
+                return x * owner.masks[self.name] if self.training else x
+
+        self.model.transformer.drop = _Mask("embed_drop")
+        for layer, block in enumerate(self.model.transformer.h):
+            block.attn.flash = False
+            block.attn.register_buffer(
+                "bias",
+                torch.tril(torch.ones(args.block_size, args.block_size)).view(
+                    1, 1, args.block_size, args.block_size
+                ),
+            )
+            block.attn.attn_dropout = _Mask(f"l{layer}_adrop")
+            block.attn.resid_dropout = _Mask(f"l{layer}_rdrop")
+            block.mlp.dropout = _Mask(f"l{layer}_mdrop")
+        with torch.no_grad():
+            for name in params:
+                for layer, target, encode in self._parts(name):
+                    source = (
+                        params[name] if layer is None else params[name][layer]
+                    )
+                    target.copy_(
+                        torch.as_tensor(encode(source), dtype=self.dtype)
+                    )
+        decay = [p for p in self.model.parameters() if p.dim() >= 2]
+        gains = [p for p in self.model.parameters() if p.dim() < 2]
+        self.optimizer = torch.optim.AdamW(
+            [
+                {"params": decay, "weight_decay": args.weight_decay},
+                {"params": gains, "weight_decay": 0.0},
+            ],
+            lr=args.learning_rate,
+            betas=(args.beta1, args.beta2),
+            eps=1e-8,
+            foreach=False,
+            fused=False,
+        )
+
+    def _parts(self, name):
+        if name in ("wte", "wpe", "ln_f"):
+            yield None, self.model.transformer[name].weight, lambda x: x
+            return
+        for layer, block in enumerate(self.model.transformer.h):
+            if name in ("wq", "wk", "wv"):
+                start = ("wq", "wk", "wv").index(name) * self.args.n_embd
+                target = block.attn.c_attn.weight[
+                    start : start + self.args.n_embd
+                ]
+                yield (
+                    layer,
+                    target,
+                    lambda x: x.reshape(self.args.n_embd, self.args.n_embd).T,
+                )
+            elif name == "wo":
+                yield (
+                    layer,
+                    block.attn.c_proj.weight,
+                    lambda x: x.reshape(self.args.n_embd, self.args.n_embd).T,
+                )
+            elif name in ("w_fc", "w_proj"):
+                attr = "c_fc" if name == "w_fc" else "c_proj"
+                yield layer, getattr(block.mlp, attr).weight, lambda x: x.T
+            else:
+                yield layer, getattr(block, name).weight, lambda x: x
+
+    def _arrays(self, name, column=None):
+        def value(parameter):
+            if column is None:
+                return parameter.detach().numpy()
+            if column == "grad":
+                return parameter.grad.detach().numpy()
+            return self.optimizer.state[parameter][column].detach().numpy()
+
+        if name in ("wte", "wpe", "ln_f"):
+            return value(self.model.transformer[name].weight)
+        result = []
+        for block in self.model.transformer.h:
+            if name in ("wq", "wk", "wv"):
+                start = ("wq", "wk", "wv").index(name) * self.args.n_embd
+                result.append(
+                    value(block.attn.c_attn.weight)[
+                        start : start + self.args.n_embd
+                    ].T.reshape(
+                        self.args.n_embd,
+                        self.args.n_head,
+                        self.args.n_embd // self.args.n_head,
+                    )
+                )
+            elif name == "wo":
+                result.append(
+                    value(block.attn.c_proj.weight).T.reshape(
+                        self.args.n_head,
+                        self.args.n_embd // self.args.n_head,
+                        self.args.n_embd,
+                    )
+                )
+            elif name in ("w_fc", "w_proj"):
+                result.append(
+                    value(
+                        getattr(
+                            block.mlp, "c_fc" if name == "w_fc" else "c_proj"
+                        ).weight
+                    ).T
+                )
+            else:
+                result.append(value(getattr(block, name).weight))
+        return np.stack(result)
+
+    def _close(self, actual, expected, label):
+        np.testing.assert_allclose(
+            actual,
+            expected,
+            atol=self.tolerance,
+            rtol=self.tolerance,
+            err_msg=label,
+        )
+
+    def evaluate(self, batch, loss):
+        self.model.eval()
+        with torch.no_grad():
+            _, expected = self.model(
+                torch.from_numpy(batch.tok.values),
+                torch.from_numpy(batch.next.values),
+            )
+        self._close(loss, float(expected), "evaluation loss")
+
+    def step(self, batch, masks, lr, loss, grads, params, moments):
+        self.model.train()
+        self.masks = {
+            name: torch.as_tensor(ds.scale.values, dtype=self.dtype)
+            for name, ds in masks.items()
+        }
+        self.optimizer.zero_grad(set_to_none=True)
+        _, expected = self.model(
+            torch.from_numpy(batch.tok.values),
+            torch.from_numpy(batch.next.values),
+        )
+        self._close(
+            loss, float(expected.detach()), "training loss with dropout"
+        )
+        expected.backward()
+        for name in params:
+            self._close(
+                grads[name].g.values,
+                self._arrays(name, "grad"),
+                f"{name} gradient",
+            )
+        if self.args.grad_clip != 0:
+            torch.nn.utils.clip_grad_norm_(
+                self.model.parameters(), self.args.grad_clip
+            )
+        for group in self.optimizer.param_groups:
+            group["lr"] = lr
+        self.optimizer.step()
+        for name in params:
+            self._close(params[name], self._arrays(name), f"{name} weight")
+            for values, col in zip(moments[name], ("exp_avg", "exp_avg_sq")):
+                self._close(values, self._arrays(name, col), f"{name} {col}")
diff --git a/tests/test_math.py b/tests/test_math.py
new file mode 100644
index 00000000..f858422c
--- /dev/null
+++ b/tests/test_math.py
@@ -0,0 +1,99 @@
+"""Error-function SQL values and analytic gradients match the array paths."""
+
+import math
+import os
+
+import numpy as np
+import pyarrow as pa
+import pytest
+import xarray as xr
+from datafusion import udf
+
+from xarray_sql import Context
+from xarray_sql.math import register_erf
+
+
+def test_erf_sql_nulls_and_replacement():
+    ctx = Context()
+    register_erf(ctx)
+    rows = ctx.sql(
+        "SELECT erf(0.0) AS zero, erf(CAST(NULL AS DOUBLE)) AS missing"
+    ).to_arrow_table()
+    assert rows.to_pydict() == {"zero": [0.0], "missing": [None]}
+    ctx.register_udf(
+        udf(
+            lambda x: pa.array(np.asarray(x) + 10),
+            [pa.float64()],
+            pa.float64(),
+            "immutable",
+            name="erf",
+        )
+    )
+    assert ctx.sql("SELECT erf(1.0) AS v").to_arrow_table().to_pydict() == {
+        "v": [11.0]
+    }
+
+
+def test_erf_replacement_is_preserved_by_dense_fallback():
+    pytest.importorskip("einfold.dense_program")
+    from xarray_sql.dense import DenseExecutor
+
+    ctx = Context()
+    register_erf(ctx)
+    ctx.from_dataset("a", xr.Dataset({"v": ("i", [0.0, 1.0])}))
+    executor = DenseExecutor(ctx)
+    run = executor.prepare("SELECT i,erf(v) AS x FROM a ORDER BY i")
+    np.testing.assert_allclose(run()["x"].to_numpy(), [0.0, math.erf(1.0)])
+    ctx.register_udf(
+        udf(
+            lambda x: pa.array(np.asarray(x) + 10),
+            [pa.float64()],
+            pa.float64(),
+            "immutable",
+            name="erf",
+        )
+    )
+    assert run().to_pydict() == {"i": [0, 1], "x": [10.0, 11.0]}
+
+
+@pytest.mark.parametrize(
+    "precision",
+    [
+        "sql",
+        "float64",
+        pytest.param(
+            "float32",
+            marks=pytest.mark.skipif(
+                os.environ.get("MLX_ENABLE_TF32") != "0",
+                reason="GPU tests require MLX_ENABLE_TF32=0",
+            ),
+        ),
+    ],
+)
+def test_exact_gelu_sql_gradient(precision):
+    ddx = pytest.importorskip("ddxdb")
+    if "erf" not in ddx.supported_functions():
+        pytest.skip("requires the optional ddx erf patch")
+    ctx = Context()
+    register_erf(ctx)
+    values = np.array([-4.0, -1.0, 0.0, 0.5, 2.0])
+    ctx.from_dataset(
+        "a",
+        xr.Dataset({"v": ("i", values)}, coords={"i": np.arange(len(values))}),
+    )
+    sql = """WITH loss AS (SELECT SUM(0.5*v*(1.0+erf(v/sqrt(2.0)))) AS l FROM a)
+        SELECT i, grad(loss,v) AS g FROM a ORDER BY i"""
+    if precision == "sql":
+        rows = ctx.sql(sql).to_arrow_table()
+    else:
+        pytest.importorskip("einfold.dense_program")
+        from xarray_sql.dense import DenseExecutor
+
+        executor = DenseExecutor(ctx, precision=precision, fallback=False)
+        rows = executor.prepare(sql)()
+    expected = 0.5 * (
+        1 + np.vectorize(math.erf)(values / np.sqrt(2))
+    ) + values * np.exp(-(values**2) / 2) / np.sqrt(2 * np.pi)
+    np.testing.assert_allclose(
+        rows["g"].to_numpy(), expected, atol=1e-6, rtol=1e-6
+    )
diff --git a/tests/test_shakespeare.py b/tests/test_shakespeare.py
new file mode 100644
index 00000000..8ef6edd8
--- /dev/null
+++ b/tests/test_shakespeare.py
@@ -0,0 +1,121 @@
+"""Real SQL training can resume without changing sampled data or updates."""
+
+import importlib.util
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+ddx = pytest.importorskip("ddxdb")
+pytest.importorskip("einfold.dense_program")
+pytestmark = pytest.mark.skipif(
+    "erf" not in ddx.supported_functions(),
+    reason="requires the optional ddx erf patch",
+)
+
+
+@pytest.mark.parametrize(
+    "execution",
+    [
+        "arrow",
+        "resident",
+        pytest.param(
+            "compiled",
+            marks=pytest.mark.skipif(
+                os.environ.get("MLX_ENABLE_TF32") != "0", reason="requires GPU"
+            ),
+        ),
+    ],
+)
+def test_shakespeare_sql_checkpoint_resume(tmp_path, monkeypatch, execution):
+    scripts = Path(__file__).resolve().parents[1] / "benchmarks/ml"
+    monkeypatch.syspath_prepend(str(scripts))
+    spec = importlib.util.spec_from_file_location(
+        "_test_shakespeare", scripts / "03_shakespeare.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    text = "To be, or not to be.\n" * 30
+    vocab = sorted(set(text))
+    ids = np.array([vocab.index(c) for c in text], dtype=np.uint16)
+    (tmp_path / "input.txt").write_text(text)
+    split = int(len(ids) * 0.9)
+    ids[:split].tofile(tmp_path / "train.bin")
+    ids[split:].tofile(tmp_path / "val.bin")
+
+    def run(out, steps, resume=None):
+        args = [
+            "03_shakespeare.py",
+            "--backend",
+            "dense_gpu" if execution == "compiled" else "dense_cpu",
+            "--data-dir",
+            str(tmp_path),
+            "--out-dir",
+            str(out),
+            "--steps",
+            str(steps),
+            "--batch_size",
+            "2",
+            "--block_size",
+            "4",
+            "--n_layer",
+            "1",
+            "--n_head",
+            "2",
+            "--n_embd",
+            "8",
+            "--eval_iters",
+            "1",
+            "--eval_interval",
+            "2",
+        ]
+        if execution == "resident":
+            args += ["--resident"]
+        elif execution == "compiled":
+            args += ["--compile-training", "--gpu-dropout"]
+        if resume:
+            args += ["--resume", str(resume)]
+        monkeypatch.setattr(sys, "argv", args)
+        module.main()
+
+    full, resumed = tmp_path / "full", tmp_path / "resumed"
+    run(full, 3)
+    run(resumed, 2)
+    run(resumed, 1, resumed / "last.npz")
+    with (
+        np.load(full / "last.npz", allow_pickle=False) as expected,
+        np.load(resumed / "last.npz", allow_pickle=False) as actual,
+    ):
+        for key in expected.files:
+            if key != "metadata":
+                if execution == "compiled":
+                    np.testing.assert_allclose(
+                        actual[key],
+                        expected[key],
+                        atol=3e-6,
+                        rtol=3e-6,
+                        err_msg=key,
+                    )
+                else:
+                    np.testing.assert_array_equal(
+                        actual[key], expected[key], err_msg=key
+                    )
+        a, b = (
+            json.loads(str(actual["metadata"])),
+            json.loads(str(expected["metadata"])),
+        )
+        assert a["iteration"] == b["iteration"] == 3
+        assert a["data_rng"] == b["data_rng"]
+        assert a["mask_rng"] == b["mask_rng"]
+        assert a["best_val_loss"] == b["best_val_loss"]
+    events = [
+        json.loads(line)
+        for line in (full / "metrics.jsonl").read_text().splitlines()
+    ]
+    assert [e["iteration"] for e in events if e["kind"] == "eval"] == [0, 2]
+    assert [e["lr"] for e in events if e["kind"] == "train"] == [
+        1e-3 * (i + 1) / 101 for i in range(3)
+    ]
diff --git a/tests/test_training.py b/tests/test_training.py
new file mode 100644
index 00000000..1f89934a
--- /dev/null
+++ b/tests/test_training.py
@@ -0,0 +1,171 @@
+"""Owned SQL training state preserves updates, failures and snapshots."""
+
+import os
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from xarray_sql import Context
+
+pytest.importorskip("ddxdb")
+pytest.importorskip("einfold.dense_program")
+from xarray_sql.training import TrainingSession
+
+_PRECISIONS = [
+    "float64",
+    pytest.param(
+        "float32",
+        marks=pytest.mark.skipif(
+            os.environ.get("MLX_ENABLE_TF32") != "0",
+            reason="GPU requires MLX_ENABLE_TF32=0",
+        ),
+    ),
+]
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_training_state_reuses_primal_and_owns_inputs(precision):
+    ctx = Context()
+    initial = xr.Dataset({"v": ("i", [1.0, 2.0, 3.0])}, coords={"i": [2, 4, 7]})
+    ctx.from_dataset("w", initial)
+    session = TrainingSession(ctx, precision=precision)
+    session.bind("w")
+    initial.v.values[:] = 99
+    step = session.prepare_step(
+        gradients={
+            "g": "WITH loss AS (SELECT SUM(v*v)/2 AS l FROM w) SELECT i,grad(loss,v) AS g FROM w ORDER BY i"
+        },
+        stages={},
+        updates={
+            "u": "SELECT p.i,p.v-0.1*g.g AS v FROM w p JOIN g ON p.i=g.i ORDER BY p.i"
+        },
+        assignments={"w": ("u", ["i", "v"])},
+    )
+    before = session.export("w")["w"]
+    for k in range(3):
+        result = step()
+        np.testing.assert_allclose(
+            float(result.loss), 7 * 0.9 ** (2 * k), rtol=1e-6
+        )
+        np.testing.assert_allclose(
+            session.export("w")["w"]["v"],
+            np.array([1, 2, 3]) * 0.9 ** (k + 1),
+            rtol=1e-6,
+        )
+    assert before.to_pydict() == {"i": [2, 4, 7], "v": [1.0, 2.0, 3.0]}
+    session.bind(
+        "w", xr.Dataset({"v": ("i", [4.0, 5.0, 6.0])}, coords={"i": [2, 4, 7]})
+    )
+    np.testing.assert_allclose(float(step().loss), 38.5, rtol=1e-6)
+    session.close()
+    with pytest.raises(RuntimeError, match="closed"):
+        step()
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_training_failure_keeps_all_state(precision):
+    ctx = Context()
+    for name, values in (("a", [1.0, 2.0]), ("b", [3.0, 4.0])):
+        ctx.from_dataset(name, xr.Dataset({"v": ("i", values)}))
+    session = TrainingSession(ctx, precision=precision)
+    session.bind("a")
+    session.bind("b")
+    before = session.export("a", "b")
+    step = session.prepare_step(
+        gradients={
+            "g": "WITH loss AS (SELECT SUM(v*v) AS l FROM a) SELECT i,grad(loss,v) AS g FROM a ORDER BY i"
+        },
+        stages={},
+        updates={
+            "good": "SELECT i,v+1 AS v FROM a ORDER BY i",
+            "bad": "SELECT i,ln(-v) AS v FROM b ORDER BY i",
+        },
+        assignments={"a": ("good", ["i", "v"]), "b": ("bad", ["i", "v"])},
+    )
+    with pytest.raises(ValueError, match="finite|exceptional"):
+        step()
+    after = session.export("a", "b")
+    assert all(after[name].equals(before[name]) for name in before)
+    assert session.version == 0
+
+
+@pytest.mark.skipif(
+    os.environ.get("MLX_ENABLE_TF32") != "0", reason="requires GPU"
+)
+def test_compiled_training_uses_new_tokens_and_learning_rates():
+    ctx = Context()
+    weights = np.arange(1, 7, dtype=float).reshape(3, 2)
+    ctx.from_dataset("w", xr.Dataset({"v": (("tok", "e"), weights)}))
+    session = TrainingSession(ctx)
+    session.bind("w")
+    session.bind("batch", xr.Dataset({"tok": ("pos", [2, 2, 0])}))
+    session.index_domains("batch", {"tok": ("w", "tok")})
+    session.bind("batch")
+    session.bind("rate", xr.Dataset({"lr": xr.DataArray(0.1)}))
+    step = session.prepare_step(
+        gradients={
+            "g": """WITH model AS (
+          SELECT b.pos,w.e,w.v FROM batch b JOIN w ON b.tok=w.tok),
+          loss AS (SELECT SUM(v*v)/2 AS l FROM model)
+          SELECT tok,e,grad(loss,v) AS g FROM w ORDER BY tok,e"""
+        },
+        stages={},
+        updates={
+            "u": "SELECT w.tok,w.e,w.v-r.lr*g.g AS v FROM w JOIN g ON w.tok=g.tok AND w.e=g.e CROSS JOIN rate r ORDER BY w.tok,w.e"
+        },
+        assignments={"w": ("u", ["tok", "e", "v"])},
+        compile=True,
+    )
+    for tokens, lr in (([2, 2, 0], 0.1), ([1, 0, 1], 0.03), ([0, 0, 2], 0.07)):
+        session.bind("batch", xr.Dataset({"tok": ("pos", tokens)}))
+        session.bind("rate", xr.Dataset({"lr": xr.DataArray(lr)}))
+        expected_loss = np.sum(weights[tokens] ** 2) / 2
+        weights = (
+            weights - lr * np.bincount(tokens, minlength=3)[:, None] * weights
+        )
+        result = step()
+        assert result.measurements["gpu_readback_bytes"] == 0
+        np.testing.assert_allclose(float(result.loss), expected_loss, rtol=3e-6)
+        rows = session.export("w")["w"]
+        np.testing.assert_allclose(
+            np.asarray(rows["v"]).reshape(3, 2), weights, rtol=3e-6
+        )
+    before = session.export("w")["w"]
+    with pytest.raises(ValueError, match="domain"):
+        session.bind("batch", xr.Dataset({"tok": ("pos", [9, 0, 1])}))
+    assert session.export("w")["w"].equals(before)
+
+
+def test_training_rejects_schema_and_coordinate_changes():
+    ctx = Context()
+    ctx.from_dataset("w", xr.Dataset({"v": ("i", [1.0, 2.0])}))
+    session = TrainingSession(ctx, precision="float64")
+    session.bind("w")
+    before = session.export("w")["w"]
+    for ds in (
+        xr.Dataset({"other": ("i", [1.0, 2.0])}),
+        xr.Dataset({"v": ("i", [1.0, 2.0])}, coords={"i": [4, 5]}),
+    ):
+        with pytest.raises(ValueError, match="schema|coordinate"):
+            session.bind("w", ds)
+    assert session.export("w")["w"].equals(before)
+
+
+@pytest.mark.parametrize("precision", _PRECISIONS)
+def test_loss_snapshot_cannot_modify_scalar_state(precision):
+    ctx = Context()
+    ctx.from_dataset("w", xr.Dataset({"v": ("i", [2.0])}))
+    session = TrainingSession(ctx, precision=precision)
+    session.bind("w")
+    step = session.prepare_step(
+        gradients={
+            "g": "WITH loss AS (SELECT SUM(v) AS l FROM w) SELECT grad(loss,v) AS g FROM w"
+        },
+        stages={},
+        updates={"u": "SELECT i,v FROM w"},
+        assignments={"w": ("u", ["i", "v"])},
+    )
+    result = step()
+    result.loss[...] = 99
+    assert session.export("w")["w"].to_pydict() == {"i": [0], "v": [2.0]}
diff --git a/xarray_sql/math.py b/xarray_sql/math.py
new file mode 100644
index 00000000..2a9c1e1d
--- /dev/null
+++ b/xarray_sql/math.py
@@ -0,0 +1,40 @@
+"""Optional SQL mathematical primitives with equivalent array lowering.
+
+DataFusion (https://datafusion.apache.org/) supplies most numerical SQL
+functions. This module supplies the error function for exact GELU, the
+activation used by nanoGPT (https://github.com/karpathy/nanoGPT).
+"""
+
+import math
+
+import numpy as np
+import pyarrow as pa
+from datafusion import udf
+
+try:
+    from scipy.special import erf as _erf
+except ImportError:
+    _erf = np.vectorize(math.erf, otypes=[np.float64])
+
+__all__ = ["register_erf"]
+
+
+def _arrow_erf(values):
+    return pa.array(
+        _erf(values.to_numpy(zero_copy_only=False)),
+        mask=values.is_null().to_numpy(zero_copy_only=False),
+        type=pa.float64(),
+    )
+
+
+def register_erf(ctx):
+    """Install SQL ``erf(DOUBLE)`` and its equivalent CPU/GPU primitive.
+
+    Differentiation also requires a ddxdb build with an erf rule; see
+    ``accelerators/ddx_math``. Replacing this UDF removes its trusted
+    primitive status so custom SQL behavior is never silently bypassed.
+    """
+    ctx.register_udf(
+        udf(_arrow_erf, [pa.float64()], pa.float64(), "immutable", name="erf")
+    )
+    ctx._trusted_math_udfs.add("erf")
diff --git a/xarray_sql/training.py b/xarray_sql/training.py
new file mode 100644
index 00000000..7cb56c03
--- /dev/null
+++ b/xarray_sql/training.py
@@ -0,0 +1,702 @@
+"""Owned array state for training programs written in SQL.
+
+EinFold (https://github.com/xqlsystems/einfold) executes the numerical
+plans and ddx (https://github.com/xqlsystems/ddx) supplies every gradient.
+This opt-in session keeps complete, nonnullable rectangular tables as
+arrays between calls. Use a dedicated Context; ordinary query APIs retain
+their independent live-Dataset and Arrow snapshot contracts.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import pyarrow as pa
+from datafusion import SessionContext
+from ddxdb import prepare as _prepare
+from einfold.dense_program import (
+    _align,
+    _axes,
+    _Axis,
+    _Column,
+    _const,
+    _Decline,
+    _expr,
+    _is,
+    _read_dataset,
+    _Sorted,
+    _Table,
+    _Tensor,
+)
+
+from xarray_sql.dense import (
+    DenseExecutor,
+    _Compiled,
+    _DenseBackend,
+    _DenseGraph,
+)
+from xarray_sql.shared import _snapshot
+
+__all__ = ["TrainingResult", "TrainingSession"]
+
+
+def _compact(table):
+    if isinstance(table, _Sorted):
+        sorts, table = table.sorts, table.table
+        if not isinstance(table, _Table):
+            raise _Decline("training requires a rectangular ordered result")
+        order = []
+        for expr in sorts:
+            column = _expr(expr.expr(), table)
+            axis = column.tensor.axis
+            if axis is None or not _is(column.valid, True):
+                raise _Decline(
+                    "training order requires nonnullable coordinate keys"
+                )
+            if axis in order:
+                continue
+            indexes = np.argsort(axis.values, kind="stable")
+            if not expr.ascending():
+                indexes = indexes[::-1]
+            if not np.array_equal(indexes, np.arange(len(axis))):
+                replacement = _Axis(axis.values[indexes])
+                table = table.take(axis, indexes, (replacement,))
+                axis = replacement
+            order.append(axis)
+        table = _Table(
+            _axes(order, table.axes), table.columns.items(), table.present
+        )
+    if not isinstance(table, _Table):
+        raise _Decline("training requires a rectangular result")
+    return table
+
+
+@dataclass(frozen=True)
+class _DeviceTable:
+    table: _Table
+    schema: pa.Schema
+
+    def select(self, names):
+        return _DeviceTable(
+            _Table(
+                self.table.axes,
+                [((None, n), self.table.resolve(None, n)) for n in names],
+                self.table.present,
+            ),
+            pa.schema([self.schema.field(n) for n in names]),
+        )
+
+
+class _TrainingBackend(_DenseBackend):
+    def returns_rows_compiled(self, plan, compiled):
+        table = _compact(compiled.run(self.owner))
+        if any(len(axis) == 0 for axis in table.axes):
+            return False
+        flags = table.present.get()
+        runtime = self.owner._runtime
+        if runtime is not None and runtime.owns(flags):
+            if self.owner._tracing:
+                raise _Decline(
+                    "compiled training needs static structural checks"
+                )
+            return runtime.boolean(flags.any())
+        return bool(np.any(flags))
+
+
+class TrainingResult:
+    """A completed update's loss and immutable device table snapshots.
+
+    Full arrays reach Arrow only through ``export``. The loss may be an
+    MLX scalar; converting it to a Python number synchronizes that scalar.
+    """
+
+    def __init__(self, session, loss, outputs, version, measurements=None):
+        self._session, self._outputs = session, outputs
+        runtime = session.executor._runtime
+        self.loss = (
+            runtime.copy(loss)
+            if runtime is not None
+            else np.array(loss, copy=True)
+        )
+        self.version = version
+        self.measurements = measurements or {}
+
+    def export(self, *names):
+        """Export independent Arrow snapshots of selected intermediate tables."""
+        return self._session._export({n: self._outputs[n] for n in names})
+
+    def loss_value(self):
+        """Read the loss scalar, recording this explicit synchronization."""
+        runtime = self._session.executor._runtime
+        if runtime is not None and runtime.owns(self.loss):
+            runtime.counts["gpu_loss_readbacks"] += 1
+            self.measurements["gpu_loss_readbacks"] = (
+                self.measurements.get("gpu_loss_readbacks", 0) + 1
+            )
+        return float(self.loss)
+
+
+class TrainingSession:
+    """Own numerical tables and execute SQL updates without Arrow intermediates.
+
+    ``bind`` imports a private copy. Later caller writes do not change it;
+    call ``bind`` again for explicit replacement. All persisted tables
+    must retain complete coordinate domains, non-NULL finite values and
+    their original schemas. An update replaces every assigned table
+    together after validation. A failed update leaves state unchanged.
+    """
+
+    def __init__(
+        self, ctx, *, precision="float32", max_cached_bytes=16 << 30, **options
+    ):
+        self.executor = DenseExecutor(
+            ctx,
+            precision=precision,
+            fallback=False,
+            max_cached_bytes=max_cached_bytes,
+            **options,
+        )
+        self.executor._backend = _TrainingBackend(self.executor)
+        self.executor._tracing = False
+        self._tables = {}
+        self._closed = False
+        self.version = 0
+        self._queries = {}
+        self._index_domains = {}
+        self._dropouts = ()
+        self._random_step = 0
+
+    def _invoke(self, function):
+        if self._closed:
+            raise RuntimeError("training session is closed")
+
+        def execute():
+            self.executor._store.update(
+                {n: t.table for n, t in self._tables.items()}
+            )
+            return function()
+
+        return self.executor._invoke(execute, None, {})
+
+    def _seal(self, table, schema):
+        table = _compact(table)
+        if len(schema) != len(table.columns):
+            raise _Decline("training result schema changed")
+        columns = []
+        for field, column in zip(schema, table.columns.values()):
+            tensor = column.tensor
+            value = tensor.get()
+            if tensor.axis is not None or tensor.host_only:
+                runtime = self.executor._runtime
+                if runtime is not None and runtime.owns(value):
+                    raise _Decline(
+                        "device-dependent columns cannot be frozen as static coordinates"
+                    )
+                frozen = _Tensor(
+                    "array",
+                    tensor.axes,
+                    value=np.asarray(value).copy(),
+                    axis=tensor.axis,
+                )
+            else:
+                frozen = _Tensor("device", tensor.axes, value=value)
+            frozen.domain = tensor.domain
+            columns.append(((None, field.name), _Column(frozen, column.valid)))
+        return _DeviceTable(_Table(table.axes, columns, table.present), schema)
+
+    def _values_and_flags(self, tables):
+        values, flags = [], []
+        for item in tables:
+            table = item.table
+            flags.append(table.present.get())
+            for col in table.columns.values():
+                flags.append(col.valid.get())
+                if col.tensor.axis is None:
+                    values.append(col.tensor.get())
+        return values, flags
+
+    def _validate(self, tables):
+        runtime = self.executor._runtime
+        values, flags = self._values_and_flags(tables)
+        if runtime is not None:
+            runtime.validate(values, flags)
+        elif not all(np.all(f) for f in flags) or not all(
+            np.isfinite(v).all() for v in values
+        ):
+            raise ValueError(
+                "training update must be rectangular, nonnullable and finite"
+            )
+
+    def bind(self, name, dataset=None):
+        """Import or explicitly replace a table, independent of caller buffers."""
+        with self.executor._lock:
+            return self._bind(name, dataset)
+
+    def _bind(self, name, dataset):
+        ctx = self.executor.ctx
+        if dataset is None:
+            dataset = ctx._registered_datasets[name]
+        private = dataset.copy(deep=True)
+        if any(size == 0 for size in private.sizes.values()):
+            raise ValueError(
+                "training tables require nonempty coordinate domains"
+            )
+        if any(
+            tuple(v.dims) != tuple(private.sizes)
+            for v in private.data_vars.values()
+        ):
+            raise ValueError(
+                "owned training columns must share the complete table dimensions"
+            )
+        if name not in ctx._registered_datasets:
+            ctx.deregister_table(name)
+            ctx.from_dataset(name, private)
+        schema = ctx.table(name).schema()
+        source = _read_dataset(name, private)
+        if schema.names != [n for _, n in source.columns] or any(
+            f.type != pa.from_numpy_dtype(c.tensor.value.dtype)
+            for f, c in zip(schema, source.columns.values())
+        ):
+            raise ValueError("replacement changes the training table schema")
+
+        def execute():
+            item = self._seal(source, schema)
+            for column, (reference, dim) in self._index_domains.get(
+                name, {}
+            ).items():
+                domain = (
+                    self._tables[reference].table.resolve(None, dim).tensor.axis
+                )
+                col = item.table.resolve(None, column)
+                values = col.tensor.get()
+                if (
+                    domain is None
+                    or values.dtype.kind not in "iu"
+                    or not np.isin(values, domain.values).all()
+                ):
+                    raise ValueError(
+                        "index values are outside the declared coordinate domain"
+                    )
+                tensor = _Tensor(
+                    "device",
+                    col.tensor.axes,
+                    value=self.executor._runtime.indices(values),
+                )
+                tensor.domain = domain
+                item.table.columns[None, column] = _Column(tensor, col.valid)
+            self._validate([item])
+            return item
+
+        item = self._invoke(execute)
+        old = self._tables.get(name)
+        if old is not None and (
+            not old.schema.equals(item.schema)
+            or len(old.table.axes) != len(item.table.axes)
+            or any(
+                not np.array_equal(a.values, b.values)
+                for a, b in zip(old.table.axes, item.table.axes)
+            )
+        ):
+            raise ValueError(
+                "replacement changes the training table schema or coordinate domains"
+            )
+        self._tables[name] = item
+
+    def index_domains(self, name, columns):
+        """Declare GPU integer input columns bounded by owned coordinate axes.
+
+        The next ``bind`` validates the actual integer values and uploads
+        them as dynamic device inputs. ``columns`` maps each field to a
+        ``(reference_table, coordinate)`` pair.
+        """
+        if self.executor._runtime is None:
+            raise ValueError(
+                "device index domains require float32 GPU execution"
+            )
+        self._index_domains[name] = dict(columns)
+        self.executor._epoch += 1
+
+    def dropout(self, names, *, probability, seed, step=0):
+        """Generate owned mask tables on GPU using an explicit seed and step.
+
+        Each named table must have one numerical field. The counter advances
+        only after a successful update. Save seed and next counter with the
+        checkpoint; identical inputs reproduce masks after resume.
+        """
+        if self.executor._runtime is None or not 0 <= probability < 1:
+            raise ValueError(
+                "GPU dropout requires float32 execution and probability in [0,1)"
+            )
+        if not 0 <= seed < 2**32 or not 0 <= step < 2**32:
+            raise ValueError("dropout seed and step must be uint32 values")
+        for name in names:
+            columns = [
+                c
+                for c in self._tables[name].table.columns.values()
+                if c.tensor.axis is None
+            ]
+            if len(columns) != 1:
+                raise ValueError(
+                    "each dropout table needs exactly one numerical field"
+                )
+        self._dropouts = tuple(sorted(names))
+        self._dropout_probability, self._random_seed, self._random_step = (
+            probability,
+            seed,
+            step,
+        )
+        self.executor._epoch += 1
+
+    def _make_masks(self, key):
+        templates = [self._tables[name] for name in self._dropouts]
+        cols = [
+            next(c for c in t.table.columns.values() if c.tensor.axis is None)
+            for t in templates
+        ]
+        arrays = self.executor._runtime.dropout(
+            key,
+            [tuple(len(a) for a in c.tensor.axes) for c in cols],
+            self._dropout_probability,
+        )
+        results = {}
+        for name, item, col, array in zip(
+            self._dropouts, templates, cols, arrays
+        ):
+            columns = [
+                (
+                    k,
+                    _Column(
+                        _Tensor("device", c.tensor.axes, value=array),
+                        _const(True),
+                    )
+                    if c is col
+                    else c,
+                )
+                for k, c in item.table.columns.items()
+            ]
+            results[name] = _DeviceTable(
+                _Table(item.table.axes, columns), item.schema
+            )
+        return results
+
+    def prepare_step(
+        self,
+        *,
+        gradients,
+        stages,
+        updates,
+        assignments,
+        loss_column="l",
+        compile=False,
+    ):
+        """Prepare one shared ddx backward followed by ordinary SQL statements.
+
+        ``gradients``, ``stages`` and ``updates`` map result names to SQL.
+        ``assignments`` maps each owned state table to ``(result, columns)``.
+        All update SELECTs read the same old state. The primal loss is read
+        from ddx's saved output, without an independent forward query.
+        """
+        step = _TrainingStep(
+            self, gradients, stages, updates, assignments, loss_column
+        )
+        step.compile = compile
+        return step
+
+    def query(self, sql):
+        """Evaluate SQL against current owned state and export its result."""
+        if sql not in self._queries:
+            self._queries[sql] = _DenseGraph(self.executor, sql)
+        graph = self._queries[sql]
+
+        def execute():
+            table = graph.run()
+            return _snapshot(table.arrow(graph.compiled[""].schema))
+
+        return self._invoke(execute)
+
+    def _export(self, tables):
+        def execute():
+            self.executor._counts["training_exports"] += len(tables)
+            return {
+                n: _snapshot(t.table.arrow(t.schema)) for n, t in tables.items()
+            }
+
+        return self._invoke(execute)
+
+    def export(self, *names):
+        """Export independent Arrow snapshots of owned tables."""
+        return self._export(
+            {n: self._tables[n] for n in (names or self._tables)}
+        )
+
+    def close(self):
+        """Release this session's owned arrays."""
+        with self.executor._lock:
+            self._tables.clear()
+            self._closed = True
+
+    def statistics(self):
+        return self.executor.statistics() | {"training_version": self.version}
+
+
+class _TrainingStep:
+    def __init__(
+        self, session, gradients, stages, updates, assignments, loss_column
+    ):
+        self.session, self.assignments, self.loss_column = (
+            session,
+            dict(assignments),
+            loss_column,
+        )
+        self.gradients, self.stages, self.updates = (
+            dict(gradients),
+            dict(stages),
+            dict(updates),
+        )
+        self.prepared = _prepare(
+            session.executor._backend, list(gradients.values())
+        )
+        self.compiled = {}
+
+    def _run(self, sql):
+        owner = self.session.executor
+        if sql not in self.compiled or self.compiled[sql].epoch != owner._epoch:
+            self.compiled[sql] = _Compiled(
+                owner, SessionContext.sql(owner.ctx, sql)
+            )
+        compiled = self.compiled[sql]
+        return self.session._seal(compiled.run(owner), compiled.schema)
+
+    def _store(self, name, item):
+        owner = self.session.executor
+        owner._store[name] = item.table
+        owner._stub(name, item.schema)
+
+    def _build(self):
+        session = self.session
+        results = {}
+        with self.prepared.run() as rewritten:
+            if not hasattr(self.prepared, "values"):
+                raise ImportError(
+                    "rebuild accelerators/ddx_math for prepared primal outputs"
+                )
+            values = [
+                v
+                for v in self.prepared.values
+                if tuple(v.columns) == (self.loss_column,)
+            ]
+            if len(values) != 1:
+                raise ValueError(
+                    "training requires one unambiguous scalar loss output"
+                )
+            loss_table = self._run(values[0].query())
+            for name, sql in zip(self.gradients, rewritten.statements):
+                results[name] = self._run(sql)
+                self._store(name, results[name])
+            for name, sql in self.stages.items():
+                results[name] = self._run(sql)
+                self._store(name, results[name])
+            for name, sql in self.updates.items():
+                results[name] = self._run(sql)
+            next_state = {
+                name: results[result].select(columns)
+                for name, (result, columns) in self.assignments.items()
+            }
+            for name, item in next_state.items():
+                old = session._tables[name]
+                if old.schema.names != item.schema.names or any(
+                    a.type != b.type for a, b in zip(old.schema, item.schema)
+                ):
+                    raise ValueError("update changes a state table schema")
+                if len(old.table.axes) != len(item.table.axes) or any(
+                    not np.array_equal(a.values, b.values)
+                    for a, b in zip(old.table.axes, item.table.axes)
+                ):
+                    raise ValueError("update changes a state coordinate domain")
+                shape = tuple(len(a) for a in item.table.axes)
+                columns = []
+                for key, col in item.table.columns.items():
+                    if col.tensor.axis is None:
+                        value = col.tensor.get()
+                        runtime = session.executor._runtime
+                        value = (
+                            runtime.broadcast(
+                                value, col.tensor.axes, item.table.axes, shape
+                            )
+                            if runtime is not None
+                            else np.broadcast_to(
+                                _align(value, col.tensor.axes, item.table.axes),
+                                shape,
+                            )
+                        )
+                        col = _Column(
+                            _Tensor("device", item.table.axes, value=value),
+                            col.valid,
+                        )
+                    columns.append((key, col))
+                next_state[name] = _DeviceTable(
+                    _Table(item.table.axes, columns, item.table.present),
+                    item.schema,
+                )
+            return loss_table, results, next_state
+
+    def _pack(self, tables):
+        arrays, layouts = [], {}
+        for name, item in tables.items():
+            columns = []
+            for key, col in item.table.columns.items():
+                t = col.tensor
+                if t.axis is None and not t.host_only:
+                    index = len(arrays)
+                    arrays.append(t.get())
+                    columns.append((key, t.axes, t.domain, index, None))
+                else:
+                    columns.append((key, t.axes, t.domain, None, col))
+            layouts[name] = (item.table.axes, item.schema, columns)
+        return arrays, layouts
+
+    @staticmethod
+    def _unpack(arrays, layouts):
+        tables = {}
+        for name, (axes, schema, columns) in layouts.items():
+            values = []
+            for key, local, domain, index, static in columns:
+                if index is None:
+                    value = static
+                else:
+                    tensor = _Tensor("device", local, value=arrays[index])
+                    tensor.domain = domain
+                    value = _Column(tensor, _const(True))
+                values.append((key, value))
+            tables[name] = _DeviceTable(_Table(axes, values), schema)
+        return tables
+
+    def _compiled_build(self):
+        session, owner = self.session, self.session.executor
+        runtime = owner._runtime
+        if runtime is None:
+            raise ValueError("compiled training requires float32 MLX execution")
+        sources = {
+            n: t
+            for n, t in sorted(session._tables.items())
+            if n not in session._dropouts
+        }
+        for item in sources.values():
+            if any(
+                c.tensor.axis is None and c.tensor.host_only
+                for c in item.table.columns.values()
+            ):
+                raise ValueError(
+                    "compiled integer inputs require explicit index_domains"
+                )
+        arrays, layouts = self._pack(sources)
+        if session._dropouts:
+            arrays.append(
+                runtime.random_key(session._random_seed, session._random_step)
+            )
+        guard = (owner._epoch, owner._backend.session_key())
+        if getattr(self, "_compiled_guard", None) != guard:
+
+            def function(*args):
+                bound = self._unpack(
+                    args[:-1] if session._dropouts else args, layouts
+                )
+                masks = (
+                    session._make_masks(args[-1]) if session._dropouts else {}
+                )
+                owner._store.clear()
+                owner._store.update({n: t.table for n, t in bound.items()})
+                owner._store.update({n: t.table for n, t in masks.items()})
+                owner._tracing = True
+                runtime.tracing = True
+                try:
+                    loss, results, next_state = self._build()
+                    results.update(masks)
+                    packed, self._output_layouts = self._pack(
+                        {
+                            ("loss",): loss,
+                            **{("result", n): t for n, t in results.items()},
+                            **{("state", n): t for n, t in next_state.items()},
+                        }
+                    )
+                    values, flags = session._values_and_flags(
+                        [loss, *results.values()]
+                    )
+                    status = runtime.status(values, flags)
+                    return packed, status
+                finally:
+                    owner._tracing = False
+                    runtime.tracing = False
+
+            self._function = runtime.compile_training(function)
+        packed, status = self._function(*arrays)
+        runtime.finish_training(packed, status)
+        restored = self._unpack(packed, self._output_layouts)
+        loss = restored[("loss",)]
+        results = {k[1]: v for k, v in restored.items() if k[0] == "result"}
+        next_state = {k[1]: v for k, v in restored.items() if k[0] == "state"}
+        owner._refresh()
+        self._compiled_guard = (owner._epoch, owner._backend.session_key())
+        return loss, results, next_state
+
+    def __call__(self):
+        session = self.session
+
+        def execute():
+            counters = session.executor.statistics()
+            if session._dropouts and session._random_step >= 2**32:
+                raise ValueError("GPU dropout counter is exhausted")
+            if getattr(self, "compile", False):
+                loss_table, results, next_state = self._compiled_build()
+            else:
+                masks = (
+                    session._make_masks(
+                        session.executor._runtime.random_key(
+                            session._random_seed, session._random_step
+                        )
+                    )
+                    if session._dropouts
+                    else {}
+                )
+                session.executor._store.update(
+                    {n: t.table for n, t in masks.items()}
+                )
+                loss_table, results, next_state = self._build()
+                results.update(masks)
+                session._validate([loss_table, *results.values()])
+            # Validation proved these full domains have no NULLs.
+            next_state = {
+                n: _DeviceTable(
+                    _Table(
+                        t.table.axes,
+                        [
+                            (k, _Column(c.tensor, _const(True)))
+                            for k, c in t.table.columns.items()
+                        ],
+                    ),
+                    t.schema,
+                )
+                for n, t in next_state.items()
+            }
+            loss = loss_table.table.resolve(None, self.loss_column).tensor.get()
+            session._tables.update(next_state)
+            session.version += 1
+            session._random_step += 1
+            session.executor._counts["training_steps"] += 1
+            measured = session.executor.statistics()
+            delta = {
+                name: measured.get(name, 0) - counters.get(name, 0)
+                for name in (
+                    "gpu_readback_bytes",
+                    "gpu_input_bytes",
+                    "gpu_status_readbacks",
+                    "training_exports",
+                    "gpu_training_compilations",
+                )
+            }
+            return TrainingResult(
+                session, loss, results, session.version, delta
+            )
+
+        return session._invoke(execute)

From 208f9b61c85abd430a1a4ba03db18034beeda103 Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:37:05 -0800
Subject: [PATCH 7/9] Describe dependency feature builds without vendored
 patches

Keep integration instructions aligned with the separate EinFold repository.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 accelerators/einfold_mlx/README.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/accelerators/einfold_mlx/README.md b/accelerators/einfold_mlx/README.md
index d82e001c..538be669 100644
--- a/accelerators/einfold_mlx/README.md
+++ b/accelerators/einfold_mlx/README.md
@@ -284,16 +284,16 @@ MLX_ENABLE_TF32=0 PYTHONPATH="$PWD" .venv/bin/python -m pytest -m 'not integrati
 MLX_ENABLE_TF32=0 .venv/bin/python -m pytest tests/test_mlx.py
 ```
 
-The native patch also passes EinFold's kernel and randomized invariant
+The native contribution also passes EinFold's kernel and randomized invariant
 suite, except for the existing plain-DataFusion `nan_compares_by_its_sign`
 assumption on this ARM Mac. That test expects `inf - inf` to produce a
 negative NaN. Here plain DataFusion produces a positive NaN, without any
-extension installed. The test remains unchanged in the patch and was
+extension installed. The test remains unchanged in the contribution and was
 documented in the measured report; it is not silently removed.
 
 ## Shared CTE execution
 
-The `0.2.0+xqlmlx.3` patch adds safe prepared-plan resets and multi-key
+The historical `0.2.0+xqlmlx.3` build added safe prepared-plan resets and multi-key
 attention detection. Batch and head keys must match in the scores, values,
 softmax partition and normalizer. The kernel validates complete positional
 domains before treating those groups as independent attention problems.

From 912730cd768b9a5b43f2d78ac8feeeace17fd941 Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:43:01 -0800
Subject: [PATCH 8/9] Pin the reviewed ddx feature branch including its report

Allow the clean persistent ddx checkout to satisfy the exact installer revision.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 accelerators/ddx_math/README.md | 2 +-
 accelerators/ddx_math/build.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/accelerators/ddx_math/README.md b/accelerators/ddx_math/README.md
index ec0eaa55..a116292b 100644
--- a/accelerators/ddx_math/README.md
+++ b/accelerators/ddx_math/README.md
@@ -9,7 +9,7 @@ the chain rule and higher derivatives; it does not substitute a model's
 handwritten backward pass.
 
 The builder uses the [ddx feature fork](https://github.com/cromwellian/ddx), pinned to commit
-`3cd5f2711221e94295f47aa6f1270f712c2edfab`, based on release commit
+`a941d46195632c9385a37e9812dd45dc10a50e94`, based on release commit
 `af71683f6678bb093e600f948c212490a4d8b74b`. Historical benchmark wheels used
 `ddxdb==0.4.6+xqlmath.3`; the contribution retains upstream version 0.4.6.
 The feature commit also exposes
diff --git a/accelerators/ddx_math/build.py b/accelerators/ddx_math/build.py
index a77c4443..a931796c 100644
--- a/accelerators/ddx_math/build.py
+++ b/accelerators/ddx_math/build.py
@@ -14,7 +14,7 @@
 import tempfile
 from pathlib import Path
 
-_REVISION = "3cd5f2711221e94295f47aa6f1270f712c2edfab"
+_REVISION = "a941d46195632c9385a37e9812dd45dc10a50e94"
 
 
 def main():

From 8982f419793508f057475a3a6e43ba34fb29b99a Mon Sep 17 00:00:00 2001
From: Ray <cromwellian@Rays-Mac-Studio.local>
Date: Sun, 1 Nov 2026 23:50:03 -0800
Subject: [PATCH 9/9] Pin the final verified dependency checkouts

Include their completed validation reports while keeping the tested numerical source unchanged.

Co-Authored-By: GPT-6 <noreply@openai.com>
---
 accelerators/ddx_math/README.md    | 2 +-
 accelerators/ddx_math/build.py     | 2 +-
 accelerators/einfold_mlx/README.md | 2 +-
 accelerators/einfold_mlx/build.py  | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/accelerators/ddx_math/README.md b/accelerators/ddx_math/README.md
index a116292b..0a892402 100644
--- a/accelerators/ddx_math/README.md
+++ b/accelerators/ddx_math/README.md
@@ -9,7 +9,7 @@ the chain rule and higher derivatives; it does not substitute a model's
 handwritten backward pass.
 
 The builder uses the [ddx feature fork](https://github.com/cromwellian/ddx), pinned to commit
-`a941d46195632c9385a37e9812dd45dc10a50e94`, based on release commit
+`953b78ecfc8bce2b1c1a71f0d4c7300444d00cc7`, based on release commit
 `af71683f6678bb093e600f948c212490a4d8b74b`. Historical benchmark wheels used
 `ddxdb==0.4.6+xqlmath.3`; the contribution retains upstream version 0.4.6.
 The feature commit also exposes
diff --git a/accelerators/ddx_math/build.py b/accelerators/ddx_math/build.py
index a931796c..a2bd6041 100644
--- a/accelerators/ddx_math/build.py
+++ b/accelerators/ddx_math/build.py
@@ -14,7 +14,7 @@
 import tempfile
 from pathlib import Path
 
-_REVISION = "a941d46195632c9385a37e9812dd45dc10a50e94"
+_REVISION = "953b78ecfc8bce2b1c1a71f0d4c7300444d00cc7"
 
 
 def main():
diff --git a/accelerators/einfold_mlx/README.md b/accelerators/einfold_mlx/README.md
index 538be669..d1406abf 100644
--- a/accelerators/einfold_mlx/README.md
+++ b/accelerators/einfold_mlx/README.md
@@ -17,7 +17,7 @@ DataFusion plans until physical planning, so ddx can serialize them through
 Substrait. MLX only evaluates arithmetic inside those plans.
 
 The released EinFold Python API has no external executor hook. The optional
-builder pins the feature commit `feaca83e4cc1c14c7cbb5f5acc734133005ad57f` in the [EinFold feature fork](https://github.com/cromwellian/einfold),
+builder pins the feature commit `9feb8e5a7efaecf6322afc436d0cac60814ef842` in the [EinFold feature fork](https://github.com/cromwellian/einfold),
 which contains the native bridge and compact array module. Its tested base
 is `39b25210d3b39027ef91584383fb114dfa8a4a48` (Python 0.2.0). Historical
 benchmark wheels were labeled `0.2.0+xqlmlx.6`; the contributed source retains
diff --git a/accelerators/einfold_mlx/build.py b/accelerators/einfold_mlx/build.py
index ba80666e..0e6b7a62 100644
--- a/accelerators/einfold_mlx/build.py
+++ b/accelerators/einfold_mlx/build.py
@@ -17,7 +17,7 @@
 import tempfile
 from pathlib import Path
 
-_REVISION = "feaca83e4cc1c14c7cbb5f5acc734133005ad57f"
+_REVISION = "9feb8e5a7efaecf6322afc436d0cac60814ef842"
 
 
 def main():