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..640e6ce7 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,19 @@ 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'. 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() + .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 +1000,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