Add an ml extra for gradients in SQL, and three ML demos

[alxmrs]

Adds an ml extra that brings the gradients of ddx, which differentiates SQL queries, into xarray_sql.Context, and three ML demos in benchmarks/ml built on it. The demos are the three target demos of einfold, a query optimizer that rewrites SQL plans into array algorithms (docs/demos.md), written for legibility rather than speed, so they can serve as benchmarks for ddx, xarray-sql and einfold.

The extra

pip install "xarray-sql[ml]" installs ddxdb>=0.3.0,<0.4. xarray_sql/ad.py follows the pattern of the geo extension:

• every Context registers ddx_stop_gradient;
• Context.sql() computes grad, vjp and jvp calls in a select list. A query that doesn't call them never reaches ddx's parser;
• xql.ad.sql_all(ctx, statements) runs several statements over one backward pass, which is how a training step updates each parameter table.

WITH pred AS (
  SELECT x.sample, SUM(x.value * w.val) AS y
  FROM x JOIN w ON x.feature = w.feature GROUP BY x.sample),
loss AS (
  SELECT AVG(power(pred.y - t.value, 2)) AS l
  FROM pred JOIN t ON pred.sample = t.sample)
SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w

The demos

Script	Model	Checked against JAX
01_attention.py	causal attention in einfold's window and GROUP BY spellings, each a view (CREATE VIEW attention AS …): the model. A separate loss reads it, as a Python loss calls its model. Plus a Nadaraya–Watson control	forward pass; gradients with respect to Q, K and V, in one statement and one backward pass; masked pairs are absent rows
02_nanogpt.py	nanoGPT's model at 1 layer and 4,416 parameters, as a view (logits) with a separate cross-entropy loss, trained on Tiny Shakespeare with SGD, or with nanoGPT's AdamW (--optimizer adamw, the learning-rate schedule as a table)	the loss, and all 11 parameter tables (and Adam's moments) after every step
03_graphcast.py	GraphCast-shaped encoder, processor and decoder on an icosahedral multi-mesh, over WeatherBench 2 ERA5	global and regional forecasts; reports the receptive field per stage

Each script runs the same SQL on every host in _hosts.py (DataFusion through xarray_sql.Context, DuckDB through xql.register). It reports time, error against JAX, and the number of operators in each host's plan. Adding a host, such as an einfold host, is one class.

Findings

The full write-up, with numbers, is in benchmarks/ml/README.md:

• Correctness: ddx's gradients match JAX to about 1e-15 everywhere it can differentiate.
• CTE recomputation: DataFusion computes a CTE once per reference. The GPT loss is 12,794 operators and 9.4 s on DataFusion, against 126 operators and 0.058 s on DuckDB (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267).
• Regional forecasts: a regional forecast takes as long as a global one on both hosts, though the region needs 28 to 54% of the grid cells.
• ddx gaps: found and filed as grad: a window aggregate (MAX OVER PARTITION BY) can't carry gradient, though its GROUP BY form can ddx#142 (window MAX), Increase cargo version (minor). #143 (GROUP BY on a dependent key) and cftime support #144 (UNION ALL, which blocks GraphCast's gradient; that row records the error).
• NaN handling: NaN is read as NULL (Document that NaN in a data variable is read as SQL NULL #268), so the mask check uses +inf.
• Stale views: a DataFusion view keeps reading a table after it is replaced (A view keeps reading a table after it is replaced: DataFusion binds views at creation #271), so the GPT recreates its model view each step.

Testing

• tests/test_ad.py: 5 tests over Datasets registered with from_dataset; the unit suite passes (553 passed).
• All three demos pass end to end on DataFusion and DuckDB. GraphCast was run at 1 and 2 processor rounds, with real ERA5.

Not done yet

• Peak memory per host.
• Part B of the NanoGPT demo (the sparsity tricks).
• The top-level README's section on gradients, held back until the feature is ready to announce.
• GraphCast's default of 2 rounds takes about 6 minutes on DataFusion.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[alxmrs]

Looking good so far! :)

[alxmrs]

Let's omit this section for now. I want to publicize it when it is really ready.

[alxmrs]

🌀 Claude: Removed in 21f99d6. The section's example still lives in benchmarks/ml/README.md and in xarray_sql.ad's docstring. I listed it under "Not done yet" in the PR description so it isn't forgotten.

[alxmrs]

In general in all documentation, if you make a reference to something, please introduce it first.

Let's put this note of feedback in the Agents.md file so future agents don't make the same mistake.

[alxmrs]

🌀 Claude: Done in 21f99d6. This docstring now says what einfold, FlashAttention (with the paper), JAX, ddx and a host are before using them. I went through every document in the PR for the same problem: the other two demos, benchmarks/ml/README.md, _hosts.py, _bench.py and xarray_sql/ad.py (Substrait, JAX), plus the GraphCast docs on #270 (GraphCast, ERA5, forcings, swish).

AGENTS.md now has the rule: introduce every project, paper, algorithm, dataset or term of art the first time a document mentions it, and link it where there is a link.

[alxmrs]

This is compact code, but I'd rather make it legible / instructive SQL code, even if it is a bit more verbose. If it's all the same perf anyway, I'd rather demo what attention looks like as a long SQL statement.

[alxmrs]

🌀 Claude: Done in 21f99d6. Each spelling is now two complete SQL statements, written out in full: the forward pass, and its gradient. The gradient statement is the forward query plus a loss CTE and a different final SELECT, so the two can be read side by side.

While doing this, I found that ddx can take the gradient with respect to all three inputs in one query, from a join of the three tables:

SELECT q.t, q.d,
       grad(loss, q.x) AS dq,
       grad(loss, k.x) AS dk,
       grad(loss, v.x) AS dv
FROM q
JOIN k ON q.t = k.t AND q.d = k.d
JOIN v ON q.t = v.t AND q.d = v.d
ORDER BY q.t, q.d

At 256 tokens it takes 0.59–0.67 s, against 0.61–0.62 s for three statements sharing one backward pass through sql_all, and it matches jax.grad to 3e-15. So the comprehension is gone, and so is sql_all in this demo. The cost is that the four CTEs appear twice per spelling.

[alxmrs]

Just want to double check: is this the same param count (and hparams) as NanoGPT? Can you add a link in a comment above this so we can double check?

[alxmrs]

In general, will you double check that our model matches the reference?

https://github.com/karpathy/nanoGPT/blob/master/model.py

[alxmrs]

🌀 Claude: Thanks for asking me to check. The hyperparameters are not nanoGPT's, on purpose. The model wasn't quite nanoGPT's either, which wasn't on purpose, and 21f99d6 fixes that.

Size. nanoGPT's Shakespeare model (config/train_shakespeare_char.py) has 6 layers, 6 heads, 384 channels and 256 characters of context: 10.65M parameters. This one has 1 layer, 2 heads, 16 channels and 16 characters: 4,416 parameters. Even at this size, the 1-layer loss is 12,794 operators and 9.4 s on DataFusion, so nanoGPT's size isn't reachable yet. Each constant now names its GPTConfig counterpart and that config's value, with links to model.py and the config, and the script prints its parameter count. The vocabulary matches: 65 characters, from sorted(set(text)), as in nanoGPT's prepare.py.

What I compared against model.py, and fixed:

• Layer-norm gains: LayerNorm has a learned weight, and the SQL had none. There are now tables ln_1, ln_2 and ln_f, with gains initialized to 1. Biases stay off, matching train.py's default bias = False.
• Activation: nanoGPT's MLP uses GELU, and mine used ReLU² (from modded-nanogpt). It's now GELU in its tanh approximation, as GPT-2 computed it. nanoGPT's nn.GELU() is the exact erf form, which DataFusion doesn't have. The two differ by less than 5e-4. ReLU² can come back as a one-line diff in Part B, where einfold's zero-skipping needs it.
• Initialization: nanoGPT scales c_proj (wo and w_proj here) to 0.02/√(2·n_layer). Mine used 0.02 for every weight; it now uses nanoGPT's scaling.

What already matched: tied wte and lm_head, learned position embeddings, pre-norm residual blocks, eps 1e-5, the 1/√head_size scale, the causal mask, and mean cross-entropy. c_attn holds the same parameters as wq, wk and wv, in one matrix.

Still different, now listed in the docstring:

• SGD at a fixed learning rate. nanoGPT uses AdamW with warmup, cosine decay, weight decay and gradient clipping.
• No dropout, so that a run can be checked against JAX. nanoGPT's Shakespeare config trains with dropout 0.2.

AdamW in SQL means keeping the optimizer's state in versioned tables, which einfold's Demo 2 calls out as an open question. I'd do it with Part B unless you want it here.

After the changes, five SGD steps still match JAX on every step: the loss to 2e-16, and all 11 parameter tables to 1e-15.

[alxmrs]

More feedback

[alxmrs]

Can we devise a way to match NanoGPT in the optimizer? For example, could we use a table of learning rates to define a LR schedule?

[alxmrs]

Maybe we should keep simple SDG and have the more complex thing behind a flag.

[alxmrs]

🌀 Claude: Done in 5f92e69, along the lines you suggested in the follow-up: SGD stays the default, and --optimizer adamw trains as nanoGPT does. The learning rate is a table, schedule(step, lr), built from nanoGPT's get_lr (linear warmup, then cosine decay to min_lr). Each step's update joins it on s.step = {step}.

An AdamW step is two phases:

1. The gradient of every parameter table, SELECT …, grad(loss, w) AS g FROM wte and so on, through sql_all, so all eleven share one backward pass. Each result is registered as the table {name}_grad.
2. One plain-SQL statement per parameter table, with no ddx in it:
   • grad_norm and clip: the global norm of all eleven gradients, through a UNION ALL of the gradient tables, and the clip factor LEAST(1.0, 1.0 / (norm + 1e-6)), as clip_grad_norm_ computes it;
   • moments: Adam's two moments, which live as columns m and v beside w in each parameter table, so the optimizer's state is versioned with the weights;
   • the update: decoupled weight decay (0 for the layer-norm gains, as in nanoGPT's configure_optimizers), bias correction from s.step + 1, and eps.

Settings: nanoGPT's Shakespeare model's: lr 1e-3 to 1e-4, β = (0.9, 0.99), weight decay 0.1, clip 1.0. One difference: the schedule is stretched to --steps (warmup is 2% of the steps, as 100 of 5,000 is), because nanoGPT's spans 5,000 steps.

Checks:

• The JAX reference step matches torch.optim.AdamW after clip_grad_norm_ exactly (difference 0.0 over four steps, with clipping active). A control that changes torch's weight decay shows a difference of 5e-4, so the comparison isn't vacuous.
• The SQL matches that reference to 6e-15 after each of five steps: the weights and both moments of all eleven tables.

Because the update is plain SQL, it would also run on DuckDB once ddx can take gradients there.

[alxmrs]

🌀 Claude: Agreed, and that's how 5f92e69 does it: --optimizer sgd (the default) or --optimizer adamw. The SGD step is still one statement per table, w - 0.5 * grad(loss, w). Details in the reply above.