feat(ddx-ad): one transpose per relational primitive; grad and vjp (M3 6)

[alxmrs]

Stacked on #73.

Rewritten after the ontology review on #80. The rules are now one per relational primitive, visible in the code. vjp takes a real cotangent, and grad requires a scalar loss, as in JAX. Nothing needs a label: a SUM is linear whatever it is summing.

One rule per primitive (transpose.rs)

JAX differentiates a function by giving each primitive a rule and composing them. A query is a composition of relational primitives, and this PR gives each one its transpose:

Primitive	SQL	Transpose
map	a projected expression y = f(x₁, …)	x̄ᵢ += ȳ · ∂f/∂xᵢ row by row, partials from ddx-core (#71)
select	WHERE, join conditions, semi-joins, a filter on a rank	the cotangent stays on the kept rows; the recomputed region contains only those
broadcast	a join	sum the cotangent back by the input's dims
reduce	grouped SUM	broadcast the group's cotangent onto every summed row

A matrix product is not a primitive. It is broadcast, map and reduce, and composing their transposes gives Ā = Σ_out C̄·B and B̄ = Σ_batch A·C̄ (design.md §4.3). COUNT's transpose is zero. Other aggregates are refused for now; MAX, MIN and AVG get rules in #76.

A saved aggregate's transpose joins the recomputed region beneath it to its cotangent on the grouping keys (IS NOT DISTINCT FROM, so a NULL group keeps its gradient). Then the region is walked right to left. Each column's cotangent is appended as a new column rather than inlined, so a shared subexpression is not duplicated.

grad and vjp (program.rs)

let program = ddx_ad::grad(&plan, &[ColumnRef::new("weight", "val")])?;

• vjp(plan, wrt) pulls a cotangent of the output back to the wrt columns. The caller supplies the cotangent as the table __ddx_cotangent, with the output's dims and its values that depend on wrt, named as in the output (BackwardProgram::cotangent lists them). The output must keep every dim, so each row is identified.
• grad(plan, wrt) is vjp of a loss seeded with 1. The output must be one row and one column, as jax.grad requires a scalar. More columns, or a row per dim, is AdError::NotScalar, and the message says to sum into one row or use vjp.

A BackwardProgram is steps to run in order, each materialized under its name:

1. __ddx_saved_{n} for each saved aggregate, then __ddx_value, the query's own result, so a program is value_and_grad;
2. __ddx_cotangent_{n} for each saved aggregate that gradient reaches, with the aggregate's dims and the cotangent of each of its values under the same names;
3. __ddx_grad_{table} for each wrt table, shaped like the table (its dims and wrt values under the table's own names), with 0 where no gradient reached.

Saved aggregates are processed parents first. A relation read in more than one place gets one contribution per read, added with UNION ALL and a grouped SUM. A join would drop the rows one contribution lacks.

Tests

ddx-datafusion/tests/grad.rs runs each program on DataFusion (with a runner in tests/common/ad.rs, which becomes public API in M4) and checks every gradient entry against a central finite difference of the query, computed by DataFusion with that one parameter perturbed. It also checks that __ddx_value reproduces the query. None of the queries has a label. Cases: an elementwise loss, a sum joined to data, two nested aggregates, a broadcast bias, a table joined to itself, SUM(…)/COUNT(…) above the aggregate, and a wrt table the loss ignores (all zeros). vjp pulls a per-group cotangent back exactly. Refusals: STDDEV, and grad of two columns or of a row per dim.

🤖 Generated with Claude Code

https://claude.ai/code/session_01CvszJhH8pn8H9G69wEPMU2

[alxmrs]

Added 915a3e8: grad now refuses a loss that may have many rows. It skipped constant inputs, so a one-row loss cross-joined to a three-row data table was accepted and the sum of three rows was differentiated. It now requires every input the output reads to be one row by the plan's own evidence (an ungrouped aggregate, or data shown to be one row). Dividing by an ungrouped COUNT(*) over data still works. Found by /code-review.

[alxmrs]

🤖😈 Adversarial review — #74 (transposes, grad, vjp)

Before the defects, credit where it's due. I ran a finite-difference sweep of 20 awkward shapes through grad at the stack tip: LayerNorm-style val - AVG(val) via a cross join, grouped variance, HAVING, IN/NOT IN semi/anti-joins, LEFT JOIN with the wrt table on the preserved side, an exp/SUM normalizer, MAX over a saved SUM, AVG of AVG, a self-join on a.i = b.i + 1, power(2.0, val), and an OR filter over a varied value. All 17 that ddx accepted match finite differences. The rule structure is sound. The three defects below are in how results are keyed and named, not in the calculus. Failing tests are pushed as 59aa505 in grad.rs.

1. Repeated dims → silently wrong gradient, then row explosion

d(i, val) = (0,1.0),(0,2.0),(1,3.0) and SUM(val*val):

gradient   : (0,6.0) (0,6.0) (1,6.0)     want (0,2.0) (0,4.0) (1,6.0)
SGD w - .25g: 5 rows out of 3 rows in

dense_gradient re-reads the table's dims, so every duplicate row gets the group's sum. The user's JOIN g ON w.i = g.i then fans out 2×2. The documented assumption is doing the damage here; see my note on #73 for a cheap run-time refusal. Test: a_table_whose_dims_repeat_is_not_given_summed_gradients.

2. The NULL-row convention is emergent, not pinned

For wn(i, val) = (0,1),(1,NULL),(2,3), the gradient at the NULL row depends on how the loss is written:

loss	grad at NULL row
SUM(val)	1.0 (the seed passes straight through)
SUM(val*val)	0.0 (2·NULL is NULL, then coalesce(…, 0))
AVG(val) (#76)	0.5

v1 fixed exactly this class of bug in #60 and pinned NULL. Design §5 requires the NULL convention to be pinned, not emergent. Separately, coalesce(grad, 0.0) in dense_gradient conflates "no contribution reached this row" with "the contribution was NULL". The first should be 0, the second probably NULL. Test: a_null_rows_gradient_does_not_depend_on_how_the_loss_is_written.

3. gradient_name isn't injective

format!("__ddx_grad_{}", names.join("_")) maps both a_b.c and a.b_c to __ddx_grad_a_b_c. The second step overwrites the first, and on the SQL surface (#83) both calls read back the second table's gradient: 2 becomes 20, with no error. Quote or escape the parts (or number the tables), and apply the same fix in sql_all's __ddx_grad_{p}_{…} and in the Python mirror. Test: tables_whose_names_join_alike_get_their_own_gradient_steps.

Smaller

• dtype: a REAL table's gradient comes back Float64, because the seeds and coalesce use f64 literals. JAX returns a gradient with its primal's dtype. Here w.val - lr * g.val silently widens a REAL parameter table to DOUBLE after one step. This is the same schema drift the nn example (feat(ddx-datafusion): nn.py trained with grad in SQL (M4 3) #80) already works around for nullability, by registering every column nullable.
• scalar_output: the one-row check is conservative, which I agree with. It refuses … GROUP BY g ORDER BY g LIMIT 1, which is fine, but the message should say it's a proof limitation rather than claiming the loss "may have many rows".

[alxmrs]

All addressed in ab69269, and the three tests in 3398f43 pass.

1. Repeated dims: each program now carries checks, which are plans that must return no rows. The dims check returns every dim tuple of a wrt table that appears more than once. ad::run runs the checks first and fails with AdError::InvalidWrt("… share their dims"). Your test asserts that refusal.
2. NULL rows: now pinned, following v1: a NULL value gets a NULL gradient, and a row nothing reached gets 0. dense_gradient is CASE WHEN val IS NULL THEN NULL ELSE coalesce(g, 0) END, so "no contribution" and "NULL value" are no longer conflated. The three spellings (SUM(val), SUM(val*val), AVG(val)) now agree.
3. Names: every table a program writes is __ddx_{id}_…, with an id that is fresh per program. Gradients are {prefix}grad_{i}_{table}, numbered by position, so a_b.c and a.b_c can't collide.

Smaller:

• dtype: the gradient is cast to its value column's type, so a REAL table gets a REAL gradient (a_gradient_has_its_values_type).
• scalar_output: the message now says ddx can't show from the plan that the query returns one row, rather than claiming it may return many.

[alxmrs]

😈🧪 Adversarial tester: ✅ approved for correctness. Reviewed at d2c7bdc (identical tree to today's rebase febfc15). Workspace tests, doctests, a 20-min soak (0 failures) and a mutation spot-check all pass. The one blocker in the stack is on #79 (#79 (comment)); the full report is in #102 (comment).

[alxmrs]

LGTM

[alxmrs]

I wonder if we could do a sort of pushdown here and omit taking the transpose of things that we never select (maybe this already happens at a higher level).