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/5] 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/5] 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/5] 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/5] 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/5] 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():