demo: use grad() inside the recursive CTE for gradient descent

Now that grad() is differentiated as a SQL rewrite before planning, it works
inside a recursive CTE — so the gradient-descent demo no longer needs the
differentiate_sql precompile step. The loss is written once and grad() is
differentiated in place within the recursion, making the whole training loop a
single declarative query with the marker in it. Updates the benchmarks README
rationale to match (the old "grad can't live in recursion via Substrait" note is
no longer true).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: alxmrs

benchmarks/README.md

@@ -34,31 +34,30 @@ Each query round-trips back to an `xarray.Dataset` via `.to_dataset(...)`.
 
 Fits a line `y ~= a*x + b` by minimising the mean squared error, with the
 **entire training loop expressed as a single recursive CTE** — no Python
-iteration. Two pieces:
-
-- **`grad` compiles the update rule.** `xql.differentiate_sql(loss, "a", cols)`
-  turns the per-row loss into its symbolic derivative *as SQL text* — the
-  autograd engine as a calculus compiler.
-- **A recursive CTE is the optimiser.** `params(step, a, b)` starts at one row
-  and each recursion appends the next generation, descending along the gradient
-  (`AVG` of the compiled rule over the data):
-
-  ```sql
-  WITH RECURSIVE params(step, a, b) AS (
-    SELECT 0, 0.0, 0.0
-    UNION ALL
-    SELECT params.step + 1, params.a - lr*AVG(da), params.b - lr*AVG(db)
-    FROM params CROSS JOIN d WHERE params.step < STEPS
-    GROUP BY params.step, params.a, params.b)
-  SELECT * FROM params ORDER BY step
-  ```
+iteration and no precompiled update rule. `grad(...)` lives *inside* the
+recursion: `params(step, a, b)` starts at one row and each recursion appends the
+next generation, descending along the gradient that `grad` computes from the
+loss formula directly (`AVG(grad(loss, a))` is the relational `d/da (Σ loss) /
+N` — differentiation through the aggregate is just linearity):
+
+```sql
+WITH RECURSIVE params(step, a, b) AS (
+  SELECT 0, 0.0, 0.0
+  UNION ALL
+  SELECT params.step + 1,
+         params.a - lr*AVG(grad(loss, a)),
+         params.b - lr*AVG(grad(loss, b))
+  FROM params CROSS JOIN d WHERE params.step < STEPS
+  GROUP BY params.step, params.a, params.b)
+SELECT * FROM params ORDER BY step
+```
 
 So gradient, update, and iteration are all declarative SQL; the trajectory is
 the rows of one query. The fit matches numpy's least-squares solution.
 Self-contained (no network).
 
-(Why differentiate to text instead of `grad(...)` inside the recursion? `grad`
-needs the Substrait round-trip, and Substrait has no recursion — so a `grad`
-marker can't live inside a recursive CTE. Differentiating once to plain SQL
-sidesteps that.)
-
+`grad` works inside the recursive CTE because it is differentiated as a SQL
+source-to-source rewrite *before* the query is planned — no Substrait round-trip,
+so no plan-shape restrictions. (If you instead want the derivative as a string to
+embed yourself, `xql.differentiate_sql(loss, "a", cols)` compiles a single
+expression to SQL text.)

benchmarks/grad_descent.py

@@ -12,29 +12,31 @@
 """Gradient descent as a single declarative SQL query.
 
 Fits a line ``y ~= a*x + b`` by minimising the mean squared error — with the
-**entire training loop expressed as one recursive CTE**, no Python iteration.
-
-Two pieces:
-
-1. **grad compiles the update rule.** ``differentiate_sql`` turns the per-row
-   loss into the symbolic derivative *as SQL text* — the autograd engine acting
-   as a calculus compiler:
-
-       da = differentiate_sql("(y-(a*x+b))^2", "a")   # -> "-2*((a*x+b)-y)*x", etc.
-
-2. **A recursive CTE is the optimiser.** ``params(step, a, b)`` starts at one
-   row and each recursion appends the next generation, descending along the
-   gradient (``AVG`` of the compiled rule over the data):
-
-       params.a - lr * AVG(da)
-
-   So the whole loop — gradient, update, and iteration — is declarative SQL;
-   the optimisation trajectory is the rows of one query.
-
-Why two pieces instead of ``grad(...)`` directly inside the recursion? ``grad``
-needs the Substrait round-trip, and Substrait has no recursion — so ``grad``
-can't live inside a recursive CTE (tracked in #194 / a follow-up). Differentiating
-once to plain SQL sidesteps that: the recursive query contains no ``grad`` marker.
+**entire training loop expressed as one recursive CTE**, no Python iteration and
+no precompiled update rule. ``grad(...)`` lives *inside* the recursion:
+
+    WITH RECURSIVE params(step, a, b) AS (
+      SELECT 0, 0.0, 0.0
+      UNION ALL
+      SELECT params.step + 1,
+             params.a - lr * AVG(grad(loss, a)),
+             params.b - lr * AVG(grad(loss, b))
+      FROM params CROSS JOIN d
+      WHERE params.step < STEPS
+      GROUP BY params.step, params.a, params.b)
+    SELECT step, a, b FROM params ORDER BY step
+
+Each recursion appends the next generation, descending along the gradient that
+``grad`` computes from the loss formula directly. ``AVG(grad(loss, a))`` is the
+relational ``d/da (Σ loss) / N`` — differentiation through the aggregate is just
+linearity. So gradient, update, and iteration are all one declarative query; the
+optimisation trajectory is the rows of that query.
+
+``grad`` is differentiated as a SQL source-to-source rewrite *before* the query
+is planned, so the marker works inside the recursive CTE (and any other query
+shape) with no Substrait round-trip. The loss is written once, as ordinary SQL,
+and the engine differentiates it symbolically — the relational equivalent of
+``jax.vmap(jax.grad(f))``, since each row is an independent evaluation point.
 
 Run standalone:
 
@@ -48,10 +50,10 @@
 
 import xarray_sql as xql
 
-# Per-row loss r^2 with residual r = y - (a*x + b), over columns a, b, x, y.
+# Per-row loss r^2 with residual r = y - (a*x + b). The columns a, b come from
+# the recursive `params` relation; x, y come from the data table `d`.
 RESIDUAL = "(y - (a * x + b))"
 LOSS = f"{RESIDUAL} * {RESIDUAL}"
-COLUMNS = ["a", "b", "x", "y"]
 LR = 0.4
 STEPS = 200
 
@@ -72,22 +74,17 @@ def main() -> None:
         chunks={"i": n},
     )
 
-    # grad compiles the per-row update rule to SQL, once.
-    da = xql.differentiate_sql(LOSS, "a", COLUMNS)
-    db = xql.differentiate_sql(LOSS, "b", COLUMNS)
-    print(f"d(loss)/da = {da}")
-    print(f"d(loss)/db = {db}\n")
-
     # The entire training loop is one declarative recursive query: each step
-    # appends the next generation, descending along the SQL-computed gradient.
+    # appends the next generation, descending along the gradient that grad()
+    # computes from the loss — differentiated inside the recursion itself.
     trajectory = ctx.sql(
         f"""
         WITH RECURSIVE params(step, a, b) AS (
           SELECT 0 AS step, CAST(0.0 AS DOUBLE) AS a, CAST(0.0 AS DOUBLE) AS b
           UNION ALL
-          SELECT params.step + 1               AS step,
-                 params.a - {LR} * AVG({da})    AS a,
-                 params.b - {LR} * AVG({db})    AS b
+          SELECT params.step + 1                        AS step,
+                 params.a - {LR} * AVG(grad({LOSS}, a))  AS a,
+                 params.b - {LR} * AVG(grad({LOSS}, b))  AS b
           FROM params CROSS JOIN d
           WHERE params.step < {STEPS}
           GROUP BY params.step, params.a, params.b
@@ -107,7 +104,8 @@ def main() -> None:
     print(f"least-squares (numpy): a={a_ols:.4f}  b={b_ols:.4f}")
     assert abs(a - a_ols) < 1e-2 and abs(b - b_ols) < 1e-2
     print(
-        "\nOK: a single recursive-CTE query fit the line to the OLS solution."
+        "\nOK: a single recursive-CTE query with grad() inside fit the line "
+        "to the OLS solution."
     )