From ec0b661a903c23e1de56ebe725876e0dada7bfac Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 28 Oct 2026 22:52:06 -0700 Subject: [PATCH 1/2] feat(ddx-ad): differentiate through UNION ALL in reverse mode grad and vjp refused a UNION ALL between a wrt table and the loss, which blocked GraphCast-shaped models: each message-passing MLP reads its inputs concatenated. Reverse mode's rules are row-local, so a union is rebuilt as one relation that holds every input's columns. Each input's rows carry its own columns, a column TRUE in them, and NULL in the others' columns, typed by a FULL JOIN to none of the others' rows. The union's columns read the input that gave the row, and their transpose sends the cotangent there and NULL to the others. Forward mode already had its rule. UNION, which removes duplicates, is still refused. Closes #144. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CvszJhH8pn8H9G69wEPMU2 --- crates/ddx-ad/src/contraction.rs | 9 + crates/ddx-ad/src/emit.rs | 21 +- crates/ddx-ad/src/forward.rs | 112 +++++++- crates/ddx-ad/src/transpose.rs | 17 +- crates/ddx-datafusion/tests/union_all.rs | 335 +++++++++++++++++++++++ docs/design.md | 18 +- python/ddxdb/tests/test_ad.py | 18 ++ python/ddxdb/tests/test_duckdb.py | 18 ++ tests/test_query_jax.py | 42 ++- 9 files changed, 580 insertions(+), 10 deletions(-) create mode 100644 crates/ddx-datafusion/tests/union_all.rs diff --git a/crates/ddx-ad/src/contraction.rs b/crates/ddx-ad/src/contraction.rs index 6a6ee1de..8acbbeb6 100644 --- a/crates/ddx-ad/src/contraction.rs +++ b/crates/ddx-ad/src/contraction.rs @@ -871,6 +871,15 @@ fn slice(region: &Region, leaf: &Leaf) -> Result> { } } Def::Const => Def::Const, + Def::Branch => Def::Branch, + Def::Union(arms) => match arms + .iter() + .map(|&(b, c)| Ok((local(b)?, local(c)?))) + .collect::>>() + { + Ok(arms) => Def::Union(arms), + Err(_) => return Ok(None), + }, }; defs.push(def); } diff --git a/crates/ddx-ad/src/emit.rs b/crates/ddx-ad/src/emit.rs index 004a1989..ae03b35e 100644 --- a/crates/ddx-ad/src/emit.rs +++ b/crates/ddx-ad/src/emit.rs @@ -5,8 +5,8 @@ //! Writing the plans of a backward program. //! //! Every step ddx emits is an ordinary Substrait plan, with nothing in it only -//! ddx understands, built from six relations: read, project, filter, join, -//! aggregate and union. The builders here produce them in the shapes +//! ddx understands, built from seven relations: read, project, filter, join, +//! aggregate, union and fetch. The builders here produce them in the shapes //! DataFusion's own producer does, since that is the consumer they are tested //! against. //! @@ -25,6 +25,7 @@ use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::aggregate_rel::{Grouping, Measure}; use substrait::proto::expression::MaskExpression; +use substrait::proto::fetch_rel::CountMode; use substrait::proto::join_rel::JoinType; use substrait::proto::plan_rel::RelType as PlanRelType; use substrait::proto::read_rel::{NamedTable, ReadType}; @@ -32,7 +33,7 @@ use substrait::proto::rel::RelType; use substrait::proto::rel_common::{Emit, EmitKind}; use substrait::proto::set_rel::SetOp; use substrait::proto::{ - r#type, AggregateFunction, AggregateRel, Expression, FilterRel, FunctionArgument, JoinRel, + r#type, AggregateFunction, AggregateRel, Expression, FetchRel, FilterRel, FunctionArgument, JoinRel, NamedStruct, Plan, PlanRel, ProjectRel, ReadRel, Rel, RelCommon, RelRoot, SetRel, }; @@ -201,6 +202,20 @@ pub fn union_all(inputs: Vec) -> Rel { } } +/// None of `input`'s rows: its columns, typed as the engine types them, +/// with nothing in them. Joined to with a FULL JOIN, it gives another +/// relation's rows NULL in those columns. +pub fn no_rows(input: Rel) -> Rel { + Rel { + rel_type: Some(RelType::Fetch(Box::new(FetchRel { + input: Some(Box::new(input)), + #[allow(deprecated)] + count_mode: Some(CountMode::Count(0)), + ..Default::default() + }))), + } +} + /// A plan whose root is `root`, with output columns named `names`. pub fn plan(mut root: Rel, names: Vec, ext: &Extensions) -> Plan { condition_joins(&mut root); diff --git a/crates/ddx-ad/src/forward.rs b/crates/ddx-ad/src/forward.rs index ed6957e4..b810fce0 100644 --- a/crates/ddx-ad/src/forward.rs +++ b/crates/ddx-ad/src/forward.rs @@ -53,9 +53,10 @@ use substrait::proto::plan_rel::RelType as PlanRelType; use substrait::proto::read_rel::ReadType; use substrait::proto::rel::RelType; use substrait::proto::rel_common::EmitKind; +use substrait::proto::set_rel::SetOp; use substrait::proto::{ AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, Expression, FetchRel, - JoinRel, NamedStruct, Plan, PlanRel, ProjectRel, ReadRel, Rel, RelCommon, SortRel, + JoinRel, NamedStruct, Plan, PlanRel, ProjectRel, ReadRel, Rel, RelCommon, SetRel, SortRel, }; use crate::elementwise::depends; @@ -158,6 +159,21 @@ pub enum Def { }, /// A column of a constant input. Const, + /// TRUE in the rows one input of a `UNION ALL` gave, NULL in the rows + /// the others gave. See [`Def::Union`]. + Branch, + /// A column of a `UNION ALL`: in the rows each input gave, that input's + /// column. Each pair is an input's [`Def::Branch`] column and its column. + /// + /// A union is rebuilt as one relation that holds every input's columns: + /// each input's rows carry its own columns and NULL in every other + /// input's. So the row-local rules apply above and beneath it, and its + /// transpose is a select: a row's cotangent goes to the column of the + /// input that gave the row, and NULL, no cotangent, to the others'. + /// Everything beneath an input is computed from its own columns, so its + /// cotangents are NULL in the rows other inputs gave, and those rows + /// send it no gradient. + Union(Vec<(usize, usize)>), } /// A rebuilt, recomputed region of row-local work. @@ -214,6 +230,12 @@ impl Region { offset: offset + shift, }, Def::Const => Def::Const, + Def::Branch => Def::Branch, + Def::Union(arms) => Def::Union( + arms.into_iter() + .map(|(b, c)| (b + shift, c + shift)) + .collect(), + ), }); } self.varied.extend(other.varied); @@ -1003,6 +1025,10 @@ impl Builder<'_> { let (s, direct) = self.cross(c)?; (s, direct, c.common.as_ref()) } + RelType::Set(u) => { + let (s, direct) = self.union_all(u)?; + (s, direct, u.common.as_ref()) + } RelType::Aggregate(a) => { let n = self.read_saved(a, below.take())?; let width = self.saved[n].outputs.len(); @@ -1305,6 +1331,90 @@ impl Builder<'_> { Ok((s, direct)) } + /// A `UNION ALL`, rebuilt as one relation that holds every input's + /// columns (see [`Def::Union`]): each input's columns, then its + /// [`Def::Branch`]; after them all, the union's columns. + fn union_all(&mut self, u: &SetRel) -> Result<(Region, Vec)> { + if SetOp::try_from(u.op) != Ok(SetOp::UnionAll) { + return Err(AdError::NotImplemented( + "a set operation other than UNION ALL on the path from a wrt table to the \ + output" + .into(), + )); + } + let mut inputs = Vec::with_capacity(u.inputs.len()); + for r in &u.inputs { + let mut s = self.lower(r)?; + s.rel = emit::project(s.rel, vec![lit_bool(true)]); + s.push(Def::Branch, false); + inputs.push(s); + } + let rels: Vec = inputs.iter().map(|s| s.rel.clone()).collect(); + let mut inputs = inputs.into_iter(); + let mut s = inputs + .next() + .ok_or_else(|| AdError::InvalidPlan("a set operation with no inputs".into()))?; + // Each input's first column here, its outputs, and its branch column. + let mut blocks = vec![(0, s.outputs.clone(), s.width() - 1)]; + for r in inputs { + let outputs = r.outputs.clone(); + let w = r.width(); + let shift = s.append(r)?; + blocks.push((shift, outputs, shift + w - 1)); + } + let width = blocks[0].1.len(); + if blocks.iter().any(|(_, outputs, _)| outputs.len() != width) { + return Err(AdError::InvalidPlan( + "a UNION ALL whose inputs have different numbers of columns".into(), + )); + } + // Each input's rows, with every other input's columns NULL: FULL + // JOINed to none of the others' rows, so the engine types the NULLs + // as it types the columns. + // + // DataFusion names a union's columns after its first input's, + // without their tables, and suffixes a repeated name: a second `x` + // is `x_1`, though a table may already have an `x_1`, and the union + // is refused. So each input's columns are renamed first, each after + // its own table and column, which no other has. + let w = s.width() + width; + let all: Vec = (0..w).collect(); + let mut rows = Vec::with_capacity(rels.len()); + for (k, (offset, outputs, _)) in blocks.iter().enumerate() { + let row = rels + .iter() + .enumerate() + .map(|(j, r)| if j == k { r.clone() } else { emit::no_rows(r.clone()) }) + .reduce(|l, r| emit::join(l, r, lit_bool(true), JoinType::Outer)) + .expect("at least one input"); + let row = emit::project(row, outputs.iter().map(|&c| field(offset + c)).collect()); + rows.push(match rels.len() { + 1 => row, + _ => rename_in_place(row, w, &all), + }); + } + s.rel = match rows.len() { + 1 => rows.pop().expect("one input"), + _ => emit::union_all(rows), + }; + let base = s.width(); + for c in 0..width { + let arms: Vec<(usize, usize)> = blocks + .iter() + .map(|(offset, outputs, branch)| (*branch, offset + outputs[c])) + .collect(); + let varied = arms.iter().any(|&(_, a)| s.varied[a]); + s.push(Def::Union(arms), varied); + } + // The union has every input's rows: one row from each is several. + if blocks.len() > 1 { + for slot in &mut s.slots { + slot.at_most_one_row = false; + } + } + Ok((s, (base..base + width).collect())) + } + /// The saved relation for aggregate `a`, read on first sight. /// /// A query that reads a CTE twice gets the CTE's plan twice: Substrait diff --git a/crates/ddx-ad/src/transpose.rs b/crates/ddx-ad/src/transpose.rs index 8c8c1430..4bea1e7d 100644 --- a/crates/ddx-ad/src/transpose.rs +++ b/crates/ddx-ad/src/transpose.rs @@ -13,6 +13,7 @@ //! | **map** | a projected expression `y = f(x₁, x₂, …)` | `x̄ᵢ += ȳ · ∂f/∂xᵢ`, row by row ([`crate::elementwise::Elementwise`]) | //! | **select** | `WHERE`, a join condition, a semi-join, a filter on a rank | the cotangent stays on the rows that were kept | //! | **broadcast** | a join, which pairs each row with every row it matches | sum the cotangent back over the rows each input row was copied to | +//! | **concatenate** | `UNION ALL` | each row's cotangent to the input that gave the row | //! | **reduce** | a grouped `SUM` | broadcast the group's cotangent to every row that was summed | //! | | `AVG` | the same, divided by the group's count | //! | | `MAX`, `MIN` | the group's cotangent to the rows that attain it, shared evenly at a tie (`jax.grad`'s convention for `jnp.max`) | @@ -313,6 +314,8 @@ impl<'a> Transposer<'a> { } } } + // A UNION ALL's inputs can give rows with the same dims. + keys.extend((0..beneath.defs.len()).filter(|&c| matches!(beneath.defs[c], Def::Branch))); Ok(keys) } @@ -650,7 +653,17 @@ impl<'a> Transposer<'a> { rel = project(rel, std::mem::take(&mut batch)); self.gather(beneath, *offset, function, &rel, here)?; } - Def::Const => { + Def::Union(arms) => { + // Each input's column gets the cotangent in the rows it + // gave, and NULL, none, in the others'. + for &(branch, col) in arms { + pending + .entry(col) + .or_default() + .push(if_then(vec![(field(branch), field(here))], null_f64())); + } + } + Def::Const | Def::Branch => { return Err(AdError::Internal(format!( "gradient reached column {c}, a {:?}, with no refusal recorded", region.defs[c] @@ -985,6 +998,8 @@ fn jitters(f: &Forward, region: &Region, col: usize) -> bool { } } Def::Window { .. } => return true, + Def::Branch => {} + Def::Union(arms) => stack.extend(arms.iter().map(|&(_, col)| col)), } } false diff --git a/crates/ddx-datafusion/tests/union_all.rs b/crates/ddx-datafusion/tests/union_all.rs new file mode 100644 index 00000000..1a9ea72c --- /dev/null +++ b/crates/ddx-datafusion/tests/union_all.rs @@ -0,0 +1,335 @@ +// SPDX-FileCopyrightText: 2026 Alexander Merose & ddx Authors +// +// SPDX-License-Identifier: Apache-2.0 + +//! `grad` through `UNION ALL`, checked against finite differences of the +//! same query. `UNION ALL` concatenates rows, so its transpose sends each +//! row's cotangent back to the input that gave the row, as the transpose of +//! `jnp.concatenate` is a slice. + +mod common; + +use common::ad::{check_gradients, rows, run, Table}; +use common::substrait_of; +use datafusion::prelude::SessionContext; +use ddx_ad::{grad, AdError, ColumnRef}; + +fn a() -> Table { + Table { + name: "a", + columns: vec![("i", "BIGINT"), ("x", "DOUBLE")], + rows: vec![vec![0.0, 1.0], vec![1.0, 2.0]], + } +} + +fn b() -> Table { + Table { + name: "b", + columns: vec![("i", "BIGINT"), ("x", "DOUBLE")], + rows: vec![vec![0.0, 3.0], vec![1.0, -0.5], vec![2.0, 0.75]], + } +} + +fn w() -> Table { + Table { + name: "w", + columns: vec![("i", "BIGINT"), ("o", "BIGINT"), ("val", "DOUBLE")], + rows: vec![ + vec![0.0, 0.0, 0.3], + vec![0.0, 1.0, -0.2], + vec![1.0, 0.0, 0.7], + vec![1.0, 1.0, 0.1], + vec![2.0, 0.0, -0.4], + vec![2.0, 1.0, 0.5], + vec![3.0, 0.0, 0.9], + vec![3.0, 1.0, -0.6], + vec![4.0, 0.0, 0.2], + vec![4.0, 1.0, 0.4], + ], + } +} + +fn ctx() -> SessionContext { + let ctx = SessionContext::new(); + ddx_datafusion::register_stop_gradient(&ctx); + ctx +} + +fn wrt(t: &str, c: &str) -> ColumnRef { + ColumnRef::new(t, c) +} + +/// `rows` in order of their keys: a step's rows come in no order. +fn sorted(rows: &[Vec]) -> Vec> { + let mut rows = rows.to_vec(); + rows.sort_by(|a, b| a.partial_cmp(b).expect("no NaN keys")); + rows +} + +#[tokio::test] +async fn a_sum_over_two_tables_concatenated() { + let grads = check_gradients( + &ctx(), + "WITH u AS (SELECT i, x FROM a UNION ALL SELECT i + 2 AS i, x FROM b) \ + SELECT SUM(x * x) AS l FROM u", + &[a(), b()], + &[wrt("a", "x"), wrt("b", "x")], + ) + .await; + // d/dx Σ x² = 2x, each row's from its own table. + assert_eq!(sorted(&grads["a"]), vec![vec![0.0, 2.0], vec![1.0, 4.0]]); +} + +#[tokio::test] +async fn one_table_read_by_both_inputs_gets_both_contributions() { + check_gradients( + &ctx(), + "WITH u AS (SELECT i, sin(x) AS v FROM a UNION ALL SELECT i, x * x AS v FROM a) \ + SELECT SUM(v * v) AS l FROM u", + &[a()], + &[wrt("a", "x")], + ) + .await; +} + +#[tokio::test] +async fn an_input_that_reads_no_wrt_table_is_concatenated_as_data() { + check_gradients( + &ctx(), + "WITH u AS (SELECT i, x FROM a UNION ALL SELECT i, x FROM b) \ + SELECT SUM(x * x * x) AS l FROM u", + &[a(), b()], + &[wrt("a", "x")], + ) + .await; +} + +#[tokio::test] +async fn groups_that_take_rows_from_both_inputs() { + // Each group of i sums one row of a and one of b, where both have it. + check_gradients( + &ctx(), + "WITH u AS (SELECT i, x FROM a UNION ALL SELECT i, 2.0 * x AS x FROM b), \ + s AS (SELECT i, SUM(x) AS t FROM u GROUP BY i) \ + SELECT SUM(t * t) AS l FROM s", + &[a(), b()], + &[wrt("a", "x"), wrt("b", "x")], + ) + .await; +} + +#[tokio::test] +async fn a_union_of_aggregates() { + check_gradients( + &ctx(), + "WITH u AS (SELECT i AS k, SUM(val) AS t FROM w GROUP BY i \ + UNION ALL SELECT o + 10 AS k, SUM(val * val) AS t FROM w GROUP BY o) \ + SELECT SUM(t * t) AS l FROM u", + &[w()], + &[wrt("w", "val")], + ) + .await; +} + +#[tokio::test] +async fn features_concatenated_then_multiplied_by_a_weight() { + // GraphCast builds an MLP's input as the UNION ALL of an edge's, its + // sender's and its receiver's features, at offset feature indices, + // and multiplies it by a weight: a contraction over the union. + let feats = Table { + name: "f", + columns: vec![("n", "BIGINT"), ("k", "BIGINT"), ("v", "DOUBLE")], + rows: vec![ + vec![0.0, 0.0, 0.5], + vec![0.0, 1.0, -1.0], + vec![1.0, 0.0, 1.5], + vec![1.0, 1.0, 0.25], + ], + }; + let more = Table { + name: "g", + columns: vec![("n", "BIGINT"), ("k", "BIGINT"), ("v", "DOUBLE")], + rows: vec![ + vec![0.0, 0.0, -0.75], + vec![0.0, 1.0, 2.0], + vec![1.0, 0.0, 0.1], + vec![1.0, 1.0, -0.3], + vec![1.0, 2.0, 0.6], + ], + }; + check_gradients( + &ctx(), + "WITH x AS (SELECT n, k, v FROM f UNION ALL SELECT n, k + 2 AS k, v FROM g), \ + h AS (SELECT x.n, w.o, SUM(x.v * w.val) AS h FROM x JOIN w ON x.k = w.i \ + GROUP BY x.n, w.o) \ + SELECT SUM(tanh(h) * tanh(h)) AS l FROM h", + &[feats, more, w()], + &[wrt("f", "v"), wrt("g", "v"), wrt("w", "val")], + ) + .await; +} + +#[tokio::test] +async fn a_maximum_tied_across_the_inputs_is_shared() { + // Both copies of the largest value attain the maximum, and share its + // gradient: one half each, so the value gets 1 in all. + let grads = check_gradients( + &ctx(), + "WITH u AS (SELECT x FROM a UNION ALL SELECT x FROM a) SELECT MAX(x) AS l FROM u", + &[a()], + &[wrt("a", "x")], + ) + .await; + assert_eq!(sorted(&grads["a"]), vec![vec![0.0, 0.0], vec![1.0, 1.0]]); +} + +#[tokio::test] +async fn a_window_over_both_inputs_rows() { + // Rows of a and b with the same i share a partition; ddx tells them + // apart by which input gave them. + check_gradients( + &ctx(), + "WITH u AS (SELECT i, x FROM a UNION ALL SELECT i, x FROM b), \ + s AS (SELECT x, SUM(x) OVER (PARTITION BY i) AS t FROM u) \ + SELECT SUM(x * t) AS l FROM s", + &[a(), b()], + &[wrt("a", "x"), wrt("b", "x")], + ) + .await; +} + +#[tokio::test] +async fn a_union_nested_in_a_union() { + check_gradients( + &ctx(), + "WITH u AS (SELECT i, x FROM a \ + UNION ALL SELECT i, x FROM (SELECT i, exp(x) AS x FROM b \ + UNION ALL SELECT i, x * x AS x FROM a) v) \ + SELECT SUM(x * x) AS l FROM u", + &[a(), b()], + &[wrt("a", "x"), wrt("b", "x")], + ) + .await; +} + +#[tokio::test] +async fn unions_of_unions_joined_over_columns_named_as_engines_rename() { + // DataFusion names a union's second `i` `i_1`; here a table already has + // one, and a union reads a union. + let c = Table { + name: "c", + columns: vec![ + ("i", "BIGINT"), + ("i_1", "BIGINT"), + ("x", "DOUBLE"), + ("x_1", "DOUBLE"), + ], + rows: vec![vec![0.0, 2.0, 3.0, 1.0], vec![1.0, 0.0, -0.5, 2.0]], + }; + check_gradients( + &ctx(), + "WITH u AS (SELECT i, x FROM c UNION ALL SELECT i, x FROM b \ + UNION ALL SELECT i_1 AS i, x_1 AS x FROM c), \ + v AS (SELECT i, x FROM (SELECT i, x FROM a UNION ALL SELECT i, x FROM b) p \ + UNION ALL SELECT i, x FROM (SELECT i, sin(x_1) AS x FROM c \ + UNION ALL SELECT i, 2.0 * x AS x FROM a) q) \ + SELECT SUM(u.x * v.x * v.x) AS l FROM u JOIN v ON u.i = v.i", + &[a(), b(), c], + &[wrt("a", "x"), wrt("b", "x"), wrt("c", "x"), wrt("c", "x_1")], + ) + .await; +} + +#[tokio::test] +async fn a_null_in_one_input_sends_its_row_no_gradient() { + let ctx = ctx(); + b().create(&ctx).await; + ctx.sql("CREATE TABLE an (i BIGINT, x DOUBLE) AS VALUES (0, 1.0), (1, NULL)") + .await + .unwrap() + .collect() + .await + .unwrap(); + let plan = substrait_of( + &ctx, + "WITH u AS (SELECT i, x FROM an UNION ALL SELECT i, x FROM b) \ + SELECT SUM(x * x) AS l FROM u", + true, + ) + .await; + let program = grad(&plan, &[wrt("an", "x"), wrt("b", "x")]).unwrap(); + run(&ctx, &program).await; + let of = |n: usize| format!("SELECT i, x FROM {} ORDER BY i", program.gradients[n].step); + let got = rows(&ctx, &of(0)).await; + assert_eq!(got[0], vec![0.0, 2.0]); + assert!( + got[1][1].is_nan(), + "the NULL row's gradient is NULL: {got:?}" + ); + assert_eq!( + rows(&ctx, &of(1)).await, + vec![vec![0.0, 6.0], vec![1.0, -1.0], vec![2.0, 1.5]] + ); +} + +#[tokio::test] +async fn a_group_whose_key_is_null_keeps_its_gradient() { + // The rows b gives the union hold NULL in the columns of s, s's key + // among them, as does s's own group whose key is NULL: those rows send + // nothing, and the group gets its own gradient. + let ctx = ctx(); + b().create(&ctx).await; + ctx.sql("CREATE TABLE n (i BIGINT, o BIGINT, x DOUBLE) AS VALUES (0, NULL, 1.0), (1, NULL, 2.0), (2, 1, 3.0)") + .await + .unwrap() + .collect() + .await + .unwrap(); + let plan = substrait_of( + &ctx, + "WITH s AS (SELECT o, SUM(x) AS t FROM n GROUP BY o), \ + u AS (SELECT o, t FROM s UNION ALL SELECT i AS o, x AS t FROM b) \ + SELECT SUM(t * t) AS l FROM u", + true, + ) + .await; + let program = grad(&plan, &[wrt("n", "x"), wrt("b", "x")]).unwrap(); + run(&ctx, &program).await; + let of = |n: usize| format!("SELECT i, x FROM {} ORDER BY i", program.gradients[n].step); + // The NULL group's t is 3 and the other's 3: d(t²)/dx = 6 for each row. + assert_eq!( + rows(&ctx, &of(0)).await, + vec![vec![0.0, 6.0], vec![1.0, 6.0], vec![2.0, 6.0]] + ); + assert_eq!( + rows(&ctx, &of(1)).await, + vec![vec![0.0, 6.0], vec![1.0, -1.0], vec![2.0, 1.5]] + ); +} + +async fn refusal(sql: &str) -> AdError { + let ctx = ctx(); + a().create(&ctx).await; + b().create(&ctx).await; + let plan = substrait_of(&ctx, sql, true).await; + grad(&plan, &[wrt("a", "x")]).expect_err("refused") +} + +#[tokio::test] +async fn a_loss_of_one_row_from_each_input_is_not_one_row() { + let err = + refusal("SELECT t FROM (SELECT SUM(x) AS t FROM a UNION ALL SELECT SUM(x) AS t FROM b) u") + .await; + assert!(matches!(err, AdError::NotScalar(_)), "{err}"); +} + +#[tokio::test] +async fn a_distinct_union_is_refused() { + // DataFusion plans UNION as a GROUP BY of every column over UNION ALL, + // so the key that depends on x is what is refused. + let err = refusal( + "WITH u AS (SELECT i, x FROM a UNION SELECT i, x FROM b) SELECT SUM(x) AS l FROM u", + ) + .await; + assert!(matches!(err, AdError::NotImplemented(_)), "{err}"); +} diff --git a/docs/design.md b/docs/design.md index 9ffff94f..b7cf9986 100644 --- a/docs/design.md +++ b/docs/design.md @@ -90,7 +90,7 @@ kept as the record of how that differentiator was designed. writes Substrait, or SQL where an engine's Substrait consumer falls short (§4.2). - Not a derivative where none is defined. An unsupported construct (a rank - used as a value, a running `MAX`, a `UNION ALL` in reverse mode for now) is + used as a value, a running `MAX`, a `UNION` that removes duplicates) is a typed error, never a silently wrong gradient (§4.6). - Not every engine at once. DataFusion first, then DuckDB-Python; a DuckDB extension, and Postgres (which needs array/XQL support first, via `pgrx`), @@ -823,6 +823,7 @@ and v2 does the same: | **map** | a projected expression `y = f(x₁, …)` | `x̄ᵢ += ȳ · ∂f/∂xᵢ`, row by row, the partials from `ddx-core`; for `CASE`, `greatest` and `least`, the derivative of the branch the row takes | | **select** | `WHERE`, a join condition, a semi-join, a filter on a rank | the cotangent stays on the rows that were kept | | **broadcast** | a join | sum the cotangent back over the rows each input row was copied to | +| **concatenate** | `UNION ALL` | each row's cotangent to the input that gave the row | | **reduce** | grouped `SUM` | broadcast the group's cotangent to every row summed | | | `AVG` | the same, divided by the group's count | | | `MAX`, `MIN` | the group's cotangent to the rows attaining it, shared evenly at a tie | @@ -851,6 +852,17 @@ a window read along several paths is transposed once, from all of their sums: a stack of windows, each layer reading the last both directly and through a window, then costs the square of its depth rather than a power of it. +**A union is rebuilt as one relation.** The rows of `UNION ALL` come from +different inputs, each with its own columns, but every rule here is row-local. +So ddx rebuilds a union as one relation that holds every input's columns: +each input's rows carry its own columns, a column that is TRUE in them, and +NULL in every other input's columns. (An input's rows are FULL JOINed to none +of the others' rows, so the engine types those NULLs as it types the +columns.) Each column of the union reads, in each row, the column of the input +that gave the row. Its transpose is a select, as `jnp.concatenate`'s is a +slice: the cotangent goes to that input's column, and NULL, no cotangent, to +the others'. So a row sends gradient only through the input that gave it. + **Map is the seam between v1 and v2.** A projected column is a scalar function of other columns of the same row, and its local derivatives are one call into `ddx-core`'s v1 `differentiate` per input column. Nothing new is built here. @@ -1173,9 +1185,7 @@ on 13 aggregates and 12 windows; a softmax written with windows equals non-spec-conformant extension-URI form DataFusion emits, and ddx keeps for the functions it adds, is untested and only matters if/when a third engine is targeted. -- `UNION ALL` in reverse mode is refused (forward mode has its rule), which - blocks the gradient of a GraphCast-shaped model in xarray-sql's ML demos. - A `GROUP BY` key computed from a `wrt` value is refused too, even when the +- A `GROUP BY` key computed from a `wrt` value is refused, even when the other keys determine it. - Windows on an engine whose Substrait consumer rejects them. ddx's rules write windows: the `AVG`, `MAX` and `MIN` rules' counts and extremes over diff --git a/python/ddxdb/tests/test_ad.py b/python/ddxdb/tests/test_ad.py index c1b11a4e..a9e84253 100644 --- a/python/ddxdb/tests/test_ad.py +++ b/python/ddxdb/tests/test_ad.py @@ -377,6 +377,24 @@ def test_not_in_over_a_subquery_is_refused(ad, ctx): + +def test_grad_through_a_union_all(ad, ctx): + # Each row of the union sends its gradient to the table it came from; + # the rows of i that both tables have land in one group. + loss = """WITH u AS (SELECT i, val FROM w UNION ALL SELECT i + 1 AS i, 2 * val AS val FROM b), + s AS (SELECT i, SUM(val) AS t FROM u GROUP BY i), + loss AS (SELECT SUM(t * t) AS l FROM s) """ + gw, gb = ad.sql_all( + ctx, + [loss + "SELECT i, grad(loss, val) FROM w", loss + "SELECT i, grad(loss, val) FROM b"], + ) + w, b = [1.0, -2.0, 0.5], [0.25, 3.0, -1.0] + t = {i: w[i] if i < 3 else 0.0 for i in range(4)} + for i in range(3): + t[i + 1] += 2 * b[i] + assert dict(pairs(gw)) == pytest.approx({i: 2 * t[i] for i in range(3)}) + assert dict(pairs(gb)) == pytest.approx({i: 4 * t[i + 1] for i in range(3)}) + def test_a_filter_on_a_gradients_dims_computes_only_those_rows(ad, ctx): # The WHERE on w's dims is pushed into the program (ddx_ad::Options::restrict); # the rows read are the unrestricted gradient's. diff --git a/python/ddxdb/tests/test_duckdb.py b/python/ddxdb/tests/test_duckdb.py index 39997e41..99d788a5 100644 --- a/python/ddxdb/tests/test_duckdb.py +++ b/python/ddxdb/tests/test_duckdb.py @@ -100,6 +100,24 @@ def test_windows_differentiate(ad, con): assert dict(pairs(rel)) == pytest.approx({0: 1 / 3 + 1.0, 1: 1 / 3, 2: 1 / 3}) + +def test_grad_through_a_union_all(ad, con): + # Each row of the union sends its gradient to the table it came from; + # the rows of i that both tables have land in one group. + loss = """WITH u AS (SELECT i, val FROM w UNION ALL SELECT i + 1 AS i, 2 * val AS val FROM b), + s AS (SELECT i, SUM(val) AS t FROM u GROUP BY i), + loss AS (SELECT SUM(t * t) AS l FROM s) """ + gw, gb = ad.sql_all( + con, + [loss + "SELECT i, grad(loss, val) FROM w", loss + "SELECT i, grad(loss, val) FROM b"], + ) + w, b = [1.0, -2.0, 0.5], [0.25, 3.0, -1.0] + t = {i: w[i] if i < 3 else 0.0 for i in range(4)} + for i in range(3): + t[i + 1] += 2 * b[i] + assert dict(pairs(gw)) == pytest.approx({i: 2 * t[i] for i in range(3)}) + assert dict(pairs(gb)) == pytest.approx({i: 4 * t[i + 1] for i in range(3)}) + def test_a_filter_on_a_gradients_dims_computes_only_those_rows(ad, con): full = pairs(ad.sql(con, GRAD)) some = pairs(ad.sql(con, GRAD + " WHERE i >= 1")) diff --git a/tests/test_query_jax.py b/tests/test_query_jax.py index a9a13340..bbabe4cc 100644 --- a/tests/test_query_jax.py +++ b/tests/test_query_jax.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 """Gradients of whole queries against `jax.grad` and `jax.jvp`: an MLP, -attention, max-pooling, and the example network's own SQL. +attention, message passing, max-pooling, and the example network's own SQL. The first three fixtures come from the small Python programs that tested ddx's design before it was built (`docs/spikes/relational_ad_spike.py`, @@ -319,6 +319,46 @@ def jax_loss(Wq, Wk, Wv, X): np.testing.assert_allclose(got[name], np.asarray(want), rtol=0, atol=TOLERANCE, err_msg=name) +def test_message_passing_matches_jax_grad(env): + jax, ad, ctx = env + jnp = jax.numpy + # GraphCast's message passing: each edge's MLP reads its own features, + # its sender's and its receiver's, concatenated: in SQL a UNION ALL at + # offset feature indices, in JAX jnp.concatenate. + rng = np.random.default_rng(3) + n_edges, n_nodes, de, dv, dh = 6, 4, 3, 2, 5 + E = rng.standard_normal((n_edges, de)) + H = rng.standard_normal((n_nodes, dv)) + W = rng.standard_normal((de + 2 * dv, dh)) * 0.5 + snd = rng.integers(0, n_nodes, n_edges) + rcv = rng.integers(0, n_nodes, n_edges) + + inputs = {"e": E, "h": H, "w": W} + register(ctx, "e", E, ("t", "d")) + register(ctx, "h", H, ("n", "d")) + register(ctx, "w", W, ("d", "o")) + edges = pa.table({"t": pa.array(range(n_edges), pa.int64()), "snd": snd, "rcv": rcv}) + ctx.register_record_batches("edges", [edges.to_batches()]) + with_loss = f""" + WITH x AS (SELECT t, d, val FROM e + UNION ALL SELECT edges.t, h.d + {de} AS d, h.val + FROM edges JOIN h ON edges.snd = h.n + UNION ALL SELECT edges.t, h.d + {de + dv} AS d, h.val + FROM edges JOIN h ON edges.rcv = h.n), + y AS (SELECT x.t, w.o, SUM(x.val * w.val) AS val + FROM x JOIN w ON x.d = w.d GROUP BY x.t, w.o), + loss AS (SELECT SUM(tanh(val) * tanh(val)) AS loss FROM y)""" + got = gradients(ad, ctx, with_loss, inputs) + + def jax_loss(E, H, W): + X = jnp.concatenate([E, H[snd], H[rcv]], axis=1) + return (jnp.tanh(X @ W) ** 2).sum() + + expected = jax.grad(jax_loss, argnums=(0, 1, 2))(*map(jnp.array, inputs.values())) + for name, want in zip(inputs, expected): + np.testing.assert_allclose(got[name], np.asarray(want), rtol=0, atol=TOLERANCE, err_msg=name) + + @pytest.mark.parametrize("tie", [False, True]) def test_max_pool_matches_jax_grad(env, tie): jax, ad, ctx = env From 045a5f706d8536b36786ff6e008c832ecf749eb9 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Thu, 29 Oct 2026 01:08:46 -0700 Subject: [PATCH 2/2] fix(ddx-ad): send a union's cotangent only to inputs that carry gradient Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CvszJhH8pn8H9G69wEPMU2 --- crates/ddx-ad/src/transpose.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/ddx-ad/src/transpose.rs b/crates/ddx-ad/src/transpose.rs index 4bea1e7d..640e6ce7 100644 --- a/crates/ddx-ad/src/transpose.rs +++ b/crates/ddx-ad/src/transpose.rs @@ -655,8 +655,10 @@ impl<'a> Transposer<'a> { } Def::Union(arms) => { // Each input's column gets the cotangent in the rows it - // gave, and NULL, none, in the others'. - for &(branch, col) in arms { + // gave, and NULL, none, in the others'. An input whose + // column reads no wrt table gets none, as a map's + // constant operand does. + for &(branch, col) in arms.iter().filter(|&&(_, col)| region.varied[col]) { pending .entry(col) .or_default()