Keep complete SQL training steps and optimizer state on GPU

.gitignore

@@ -17,3 +17,10 @@ CHANGELOG.md
 .direnv
 # Pyodide wheels copied in to test the browser demo locally
 docs/demo/*.whl
+# SQL training checkpoints; retain the accompanying text metrics/logs.
+out-shakespeare-sql/
+benchmarks/apple/shakespeare-matched-*/last.npz
+benchmarks/apple/shakespeare-matched-*/best.npz
+benchmarks/apple/shakespeare-resident-*/last.npz
+benchmarks/apple/shakespeare-compiled-*/last.npz
+benchmarks/apple/shakespeare-compiled-*/best.npz

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 }
 

README.md

@@ -289,6 +289,26 @@ I want to give a special thanks to the following folks and institutions:
   opinion.
 
 
+## Apple silicon tensor acceleration
+
+The optional [Apple silicon integration](accelerators/einfold_mlx/README.md)
+connects [EinFold](https://github.com/xqlsystems/einfold), a tensor optimizer
+for SQL, to [MLX](https://github.com/ml-explore/mlx), Apple's array framework.
+It keeps SQL planning and gradients from
+[ddx](https://github.com/xqlsystems/ddx), which differentiates SQL queries.
+Float64 CPU execution is the default;
+float32 GPU execution requires explicit opt-in. See the
+[local M5 Ultra measurements](benchmarks/apple/README.md) for end-to-end
+results, numerical errors, dispatch thresholds and limitations.
+The [shared-CTE nanoGPT comparison](benchmarks/apple/nanogpt-shared.md)
+combines cached shared SQL stages with the same CPU/GPU kernels and checks
+larger models against JAX.
+The experimental [compact execution path](benchmarks/apple/dense-execution.md)
+retains arrays across forward and backward stages, reuses plans, and
+executes supported operators through MLX. Its report includes checked
+five-step GPT-2-size training, memory bounds and remaining native-framework
+performance gaps.
+
 ## License
 
 ```

accelerators/ddx_math/README.md

@@ -0,0 +1,44 @@
+# Exact GELU for SQL training
+
+[ddx](https://github.com/xqlsystems/ddx) differentiates SQL query plans.
+Its released scalar rules do not include the error function, `erf`, used
+by [nanoGPT](https://github.com/karpathy/nanoGPT)'s exact GELU activation.
+The pinned feature commit adds `erf'(x) = 2/sqrt(pi) * exp(-x*x)` to the scalar
+registry and retains `erf` in derivative plans. The rule composes with
+the chain rule and higher derivatives; it does not substitute a model's
+handwritten backward pass.
+
+The builder uses the [ddx feature fork](https://github.com/cromwellian/ddx), pinned to commit
+`953b78ecfc8bce2b1c1a71f0d4c7300444d00cc7`, based on release commit
+`af71683f6678bb093e600f948c212490a4d8b74b`. Historical benchmark wheels used
+`ddxdb==0.4.6+xqlmath.3`; the contribution retains upstream version 0.4.6.
+The feature commit also exposes
+`Prepared.values`: metadata for the primal outputs computed by prepared
+derivative programs. These outputs can be read during a run before its
+temporary tables are released. This lets the GPU training session reuse
+the actual loss from backward preparation instead of executing a separate
+forward query. It does not alter existing derivative rules.
+
+Engines must
+register an `erf(DOUBLE)` implementation; `xarray_sql.math.register_erf`
+does this for DataFusion and marks its equivalent array primitive as safe
+for EinFold CPU and MLX GPU execution. Replacing the SQL UDF removes that
+permission. This contribution does not add an erf implementation to DuckDB.
+
+From the xarray-sql directory, after installing the normal ML dependencies
+and the [compact EinFold executor](../einfold_mlx/README.md):
+
+```bash
+source .venv/bin/activate
+uv pip install cmake maturin
+python accelerators/ddx_math/build.py
+```
+
+Requires git, a Rust toolchain, CMake and uv on PATH. The builder clones
+an isolated checkout, validates the pinned feature revision, builds its
+locked wheel and installs it in the active virtual environment. No patch
+is applied. Use `--source /path/to/ddx` for a clean checkout already at the pin. `uv sync` may
+replace locally built wheels; rerun the optional builders if necessary.
+
+The public SQL numerical checks are in `tests/test_math.py`. The ddx contribution
+also adds a Rust check of the chain rule and second derivative.

accelerators/ddx_math/build.py

@@ -0,0 +1,90 @@
+"""Build ddx's exact error-function derivative in an isolated checkout.
+
+ddx (https://github.com/xqlsystems/ddx) differentiates SQL. The pinned
+feature commit adds erf's analytic derivative, including higher derivatives, and
+exposes prepared primal output metadata for loss reuse.
+Requires git, cargo, CMake, maturin and uv on PATH.
+"""
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+_REVISION = "953b78ecfc8bce2b1c1a71f0d4c7300444d00cc7"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--source", type=Path, help="reuse a clean pinned checkout"
+    )
+    args = parser.parse_args()
+    if sys.prefix == sys.base_prefix:
+        parser.error("use the repository virtual environment's Python")
+    for command in ("git", "cargo", "cmake", "uv"):
+        if shutil.which(command) is None:
+            parser.error(f"{command} must be on PATH")
+    with tempfile.TemporaryDirectory(prefix="xql-ddx-math-") as temporary:
+        source = args.source or Path(temporary) / "source"
+        if args.source is None:
+            subprocess.run(
+                [
+                    "git",
+                    "clone",
+                    "https://github.com/cromwellian/ddx.git",
+                    str(source),
+                ],
+                check=True,
+            )
+            subprocess.run(
+                ["git", "checkout", _REVISION], cwd=source, check=True
+            )
+        if (
+            subprocess.check_output(
+                ["git", "rev-parse", "HEAD"], cwd=source, text=True
+            ).strip()
+            != _REVISION
+        ):
+            parser.error(f"expected ddx revision {_REVISION}")
+        if subprocess.check_output(
+            ["git", "status", "--porcelain"], cwd=source
+        ):
+            parser.error("source checkout must be clean")
+        wheels = Path(temporary) / "wheels"
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "maturin",
+                "build",
+                "--release",
+                "--locked",
+                "--out",
+                str(wheels),
+            ],
+            cwd=source / "python/ddxdb",
+            env=os.environ | {"PYO3_PYTHON": sys.executable},
+            check=True,
+        )
+        (wheel,) = wheels.glob("*.whl")
+        subprocess.run(
+            [
+                "uv",
+                "pip",
+                "install",
+                "--python",
+                sys.executable,
+                "--no-deps",
+                "--reinstall",
+                str(wheel),
+            ],
+            check=True,
+        )
+
+
+if __name__ == "__main__":
+    main()