From 1512fe94cfefbee1b9d76b1f93fd48fd00121aa4 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Tue, 7 Jul 2026 15:44:05 +0900 Subject: [PATCH 1/2] classify linear transpose inputs --- examples/eager_reverse_mode.rs | 16 +- examples/gradient_two_inputs.rs | 14 +- examples/primitive_linearization.rs | 16 +- src/eager/record.rs | 8 +- src/lib.rs | 4 +- src/linear_transpose.rs | 127 ++++++++++---- src/linearize.rs | 85 ++++++++- src/linearized_graph.rs | 23 ++- src/rules/mod.rs | 3 +- src/rules/primitive_builder.rs | 71 +++++--- src/rules/primitive_op.rs | 10 +- tests/adcontext_tests.rs | 6 +- tests/common/linearize_macros.rs | 2 +- tests/common/mod.rs | 6 +- tests/common/transpose_macros.rs | 14 +- tests/complex_ad_tests.rs | 6 +- tests/eager_backward_tests.rs | 14 +- tests/eager_record_tests.rs | 12 +- tests/edge_case_tests.rs | 10 +- tests/fallible_ad_tests.rs | 4 +- tests/robustness_tests.rs | 8 +- tests/rules_public_api_tests.rs | 4 +- tests/transpose_input_classification_tests.rs | 163 ++++++++++++++++++ tests/vector_ad_tests.rs | 6 +- 24 files changed, 509 insertions(+), 123 deletions(-) create mode 100644 tests/transpose_input_classification_tests.rs diff --git a/examples/eager_reverse_mode.rs b/examples/eager_reverse_mode.rs index 0be19d0..b2cf47a 100644 --- a/examples/eager_reverse_mode.rs +++ b/examples/eager_reverse_mode.rs @@ -6,7 +6,7 @@ use computegraph::{EvaluableGraphOperation, GraphOperation}; use tidu::eager::{self, BackwardExecutor, EagerInput, KeySource, RecordedGraph, Recorder}; use tidu::{ linear_transpose_with_builder, ADKey, ADRuleResult, DiffPassId, LinearizedGraph, Primitive, - PrimitiveBuilder, PrimitiveGraph, PrimitiveValue, + PrimitiveBuilder, PrimitiveGraph, PrimitiveTransposeInput, PrimitiveValue, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -151,7 +151,7 @@ impl Primitive for ScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_outputs: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -390,7 +390,7 @@ fn sum_tangent_terms( fn transpose_mul( builder: &mut impl PrimitiveBuilder, - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], ct: LocalValueId, role: &OperationRole, ) -> Vec> { @@ -402,7 +402,10 @@ fn transpose_mul( if active_mask[0] { let out = builder.add_primitive( ScalarOp::Mul, - vec![inputs[1].clone(), PrimitiveValue::Local(ct)], + vec![ + inputs[1].as_residual_value().unwrap(), + PrimitiveValue::Local(ct), + ], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -412,7 +415,10 @@ fn transpose_mul( if active_mask[1] { let out = builder.add_primitive( ScalarOp::Mul, - vec![inputs[0].clone(), PrimitiveValue::Local(ct)], + vec![ + inputs[0].as_residual_value().unwrap(), + PrimitiveValue::Local(ct), + ], OperationRole::Linearized { active_mask: vec![false, true], }, diff --git a/examples/gradient_two_inputs.rs b/examples/gradient_two_inputs.rs index 74ae459..46a5877 100644 --- a/examples/gradient_two_inputs.rs +++ b/examples/gradient_two_inputs.rs @@ -15,7 +15,7 @@ use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef}; use computegraph::{EvaluableGraphOperation, GraphOperation}; use tidu::{ linear_transpose, linearize, ADKey, DiffPassId, LinearizedGraph, Primitive, PrimitiveBuilder, - PrimitiveValue, + PrimitiveTransposeInput, PrimitiveValue, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -122,7 +122,7 @@ impl Primitive for ScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_outputs: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -141,7 +141,10 @@ impl Primitive for ScalarOp { if active_mask[0] { let out = builder.add_primitive( Self::Mul, - vec![inputs[1].clone(), PrimitiveValue::Local(ct)], + vec![ + inputs[1].as_residual_value().unwrap(), + PrimitiveValue::Local(ct), + ], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -151,7 +154,10 @@ impl Primitive for ScalarOp { if active_mask[1] { let out = builder.add_primitive( Self::Mul, - vec![inputs[0].clone(), PrimitiveValue::Local(ct)], + vec![ + inputs[0].as_residual_value().unwrap(), + PrimitiveValue::Local(ct), + ], OperationRole::Linearized { active_mask: vec![false, true], }, diff --git a/examples/primitive_linearization.rs b/examples/primitive_linearization.rs index 922585f..b03bb17 100644 --- a/examples/primitive_linearization.rs +++ b/examples/primitive_linearization.rs @@ -9,7 +9,7 @@ use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef}; use computegraph::{EvaluableGraphOperation, GraphOperation}; use tidu::{ linear_transpose, linearize, ADKey, DiffPassId, LinearizedGraph, Primitive, PrimitiveBuilder, - PrimitiveValue, + PrimitiveTransposeInput, PrimitiveValue, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -154,7 +154,7 @@ impl Primitive for ScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_outputs: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -216,7 +216,7 @@ fn sum_tangent_terms( fn transpose_mul( builder: &mut impl PrimitiveBuilder, - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], ct: LocalValueId, role: &OperationRole, ) -> Vec> { @@ -228,7 +228,10 @@ fn transpose_mul( if active_mask[0] { let out = builder.add_primitive( ScalarOp::Mul, - vec![inputs[1].clone(), PrimitiveValue::Local(ct)], + vec![ + inputs[1].as_residual_value().unwrap(), + PrimitiveValue::Local(ct), + ], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -238,7 +241,10 @@ fn transpose_mul( if active_mask[1] { let out = builder.add_primitive( ScalarOp::Mul, - vec![inputs[0].clone(), PrimitiveValue::Local(ct)], + vec![ + inputs[0].as_residual_value().unwrap(), + PrimitiveValue::Local(ct), + ], OperationRole::Linearized { active_mask: vec![false, true], }, diff --git a/src/eager/record.rs b/src/eager/record.rs index cfc048b..fcf4fc4 100644 --- a/src/eager/record.rs +++ b/src/eager/record.rs @@ -404,7 +404,7 @@ impl Recorder { /// use std::sync::Arc; /// use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey}; /// use tidu::{ - /// ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue, + /// ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, /// }; /// use tidu::eager::{EagerInput, KeySource, RecordedGraph, Recorder}; /// @@ -452,7 +452,7 @@ impl Recorder { /// &self, /// _builder: &mut impl PrimitiveBuilder, /// cotangent_out: &[Option], - /// _inputs: &[PrimitiveValue], + /// _inputs: &[PrimitiveTransposeInput], /// _role: &OperationRole, /// _ctx: &mut (), /// ) -> tidu::ADRuleResult>> { @@ -582,7 +582,7 @@ fn fresh_value_keys( #[cfg(test)] mod tests { use super::*; - use crate::{DiffPassId, PrimitiveBuilder, PrimitiveValue}; + use crate::{DiffPassId, PrimitiveBuilder, PrimitiveTransposeInput}; use computegraph::LocalValueId; #[derive(Clone, Debug, Hash, PartialEq, Eq)] @@ -641,7 +641,7 @@ mod tests { &self, _builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - _inputs: &[PrimitiveValue], + _inputs: &[PrimitiveTransposeInput], _role: &OperationRole, _ctx: &mut (), ) -> ADRuleResult>> { diff --git a/src/lib.rs b/src/lib.rs index 429245c..ba75f6e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -44,9 +44,9 @@ pub mod rules; pub use linear_transpose::{linear_transpose, linear_transpose_with_builder}; pub use linearize::linearize; -pub use linearized_graph::LinearizedGraph; +pub use linearized_graph::{LinearizedGraph, LinearizedValueKind}; pub use primitive_graph::PrimitiveGraph; pub use rules::{ ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive, PrimitiveBuilder, - PrimitiveValue, + PrimitiveTransposeInput, PrimitiveValue, }; diff --git a/src/linear_transpose.rs b/src/linear_transpose.rs index 8b41831..a4012c8 100644 --- a/src/linear_transpose.rs +++ b/src/linear_transpose.rs @@ -1,8 +1,8 @@ use std::collections::HashMap; -use crate::rules::GraphPrimitiveBuilder; use crate::{ - ADKey, ADRuleError, ADRuleKind, ADRuleResult, Primitive, PrimitiveBuilder, PrimitiveValue, + ADKey, ADRuleError, ADRuleKind, ADRuleResult, LinearizedValueKind, Primitive, PrimitiveBuilder, + PrimitiveTransposeInput, PrimitiveValue, }; use computegraph::graph::GraphBuilder; use computegraph::{LocalValueId, OperationRole, ValueKey, ValueRef}; @@ -33,6 +33,7 @@ where let mut cotangent_env: HashMap, LocalValueId> = HashMap::new(); let mut cotangent_seed_inputs = Vec::new(); let graph = linear.as_graph(); + let mut value_kinds = Vec::new(); for (index, maybe_tangent_output) in linear.tangent_outputs().iter().enumerate() { let tangent_output_id = match maybe_tangent_output { @@ -43,6 +44,8 @@ where let source_key = graph.values()[*tangent_output_id].key.clone(); let seed_key = cotangent_seed_key(linear, index)?; let seed_id = builder.add_input(seed_key.clone()); + debug_assert_eq!(seed_id, value_kinds.len()); + value_kinds.push(LinearizedValueKind::Linear { primal: None }); cotangent_env.insert(source_key, seed_id); cotangent_seed_inputs.push((seed_key, seed_id)); } @@ -57,18 +60,13 @@ where continue; } - let rule_inputs: Vec> = op_node + let rule_inputs: Vec> = op_node .inputs .iter() - .map(|input| match input { - ValueRef::Local(local_id) => { - PrimitiveValue::External(graph.values()[*local_id].key.clone()) - } - ValueRef::External(key) => PrimitiveValue::External(key.clone()), - }) + .map(|input| transpose_input_for(linear, input)) .collect(); - let mut primitive_builder = GraphPrimitiveBuilder::new(&mut builder); + let mut primitive_builder = TrackingPrimitiveBuilder::new(&mut builder, &mut value_kinds); let cotangent_in = op_node.operation.transpose_rule( &mut primitive_builder, &cotangent_out, @@ -93,16 +91,12 @@ where Some(cotangent_id) => cotangent_id, None => continue, }; - let input_key = match input { - PrimitiveValue::Local(_) => { - unreachable!("rule inputs are normalized to external refs") - } - PrimitiveValue::External(key) => key.clone(), - }; + let input_key = input.key().clone(); match cotangent_env.get(&input_key).copied() { Some(existing_id) => { - let mut primitive_builder = GraphPrimitiveBuilder::new(&mut builder); + let mut primitive_builder = + TrackingPrimitiveBuilder::new(&mut builder, &mut value_kinds); let sum = primitive_builder.add_primitive( Op::add(), vec![ @@ -139,6 +133,7 @@ where builder.build(), cotangent_seed_inputs, tangent_outputs, + value_kinds, )) } @@ -174,15 +169,10 @@ where continue; } - let rule_inputs: Vec> = op_node + let rule_inputs: Vec> = op_node .inputs .iter() - .map(|input| match input { - ValueRef::Local(local_id) => { - PrimitiveValue::External(graph.values()[*local_id].key.clone()) - } - ValueRef::External(key) => PrimitiveValue::External(key.clone()), - }) + .map(|input| transpose_input_for(linear, input)) .collect(); let cotangent_in = op_node.operation.transpose_rule( @@ -209,12 +199,7 @@ where Some(cotangent_id) => cotangent_id, None => continue, }; - let input_key = match input { - PrimitiveValue::Local(_) => { - unreachable!("rule inputs are normalized to external refs") - } - PrimitiveValue::External(key) => key.clone(), - }; + let input_key = input.key().clone(); match cotangent_env.get(&input_key).copied() { Some(existing_id) => { @@ -247,6 +232,88 @@ where .collect()) } +fn transpose_input_for( + linear: &LinearizedGraph, + input: &ValueRef, +) -> PrimitiveTransposeInput +where + Op::InputKey: ADKey, +{ + let graph = linear.as_graph(); + match input { + ValueRef::Local(local_id) => { + let key = graph.values()[*local_id].key.clone(); + match linear.value_kind(*local_id) { + LinearizedValueKind::Residual => PrimitiveTransposeInput::Residual(key), + LinearizedValueKind::Linear { primal } => PrimitiveTransposeInput::Linear { + key, + primal: primal.clone(), + }, + } + } + ValueRef::External(key) => PrimitiveTransposeInput::Residual(key.clone()), + } +} + +struct TrackingPrimitiveBuilder<'a, Op: computegraph::GraphOperation> { + inner: &'a mut GraphBuilder, + value_kinds: &'a mut Vec>, +} + +impl<'a, Op: computegraph::GraphOperation> TrackingPrimitiveBuilder<'a, Op> { + fn new( + inner: &'a mut GraphBuilder, + value_kinds: &'a mut Vec>, + ) -> Self { + Self { inner, value_kinds } + } + + fn output_kind( + &self, + inputs: &[PrimitiveValue], + role: &OperationRole, + ) -> LinearizedValueKind { + let OperationRole::Linearized { active_mask } = role else { + return LinearizedValueKind::Residual; + }; + + if inputs.iter().enumerate().any(|(index, input)| { + active_mask.get(index).copied().unwrap_or(false) && self.is_linear_input(input) + }) { + LinearizedValueKind::Linear { primal: None } + } else { + LinearizedValueKind::Residual + } + } + + fn is_linear_input(&self, input: &PrimitiveValue) -> bool { + match input { + PrimitiveValue::Local(id) => { + matches!(self.value_kinds[*id], LinearizedValueKind::Linear { .. }) + } + PrimitiveValue::External(_) => true, + } + } +} + +impl PrimitiveBuilder for TrackingPrimitiveBuilder<'_, Op> { + fn add_primitive( + &mut self, + op: Op, + inputs: Vec>, + role: OperationRole, + ) -> Vec { + let output_kind = self.output_kind(&inputs, &role); + let inputs = inputs.into_iter().map(ValueRef::from).collect(); + let outputs = self.inner.add_operation(op, inputs, role); + for output_id in &outputs { + debug_assert_eq!(*output_id, self.value_kinds.len()); + self.value_kinds.push(output_kind.clone()); + } + outputs + } +} + fn cotangent_seed_key( linear: &LinearizedGraph, index: usize, diff --git a/src/linearize.rs b/src/linearize.rs index 96434b0..0d0b5e5 100644 --- a/src/linearize.rs +++ b/src/linearize.rs @@ -1,11 +1,13 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use crate::rules::GraphPrimitiveBuilder; -use crate::{ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive}; +use crate::rules::{PrimitiveBuilder, PrimitiveValue}; +use crate::{ + ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, LinearizedValueKind, Primitive, +}; use computegraph::graph::GraphBuilder; use computegraph::resolve::{ResolvedView, ValueDef}; -use computegraph::{GraphOperation, LocalValueId, OperationKey, ValueKey}; +use computegraph::{GraphOperation, LocalValueId, OperationKey, OperationRole, ValueKey, ValueRef}; use crate::LinearizedGraph; @@ -43,11 +45,16 @@ where let topo_keys = topological_order(view, outputs, aliases); let mut tangent_env: HashMap, Option> = HashMap::new(); let mut processed_ops = HashSet::new(); + let mut value_kinds = Vec::new(); let mut tangent_inputs = Vec::with_capacity(wrt.len()); for wrt_key in wrt { let tangent_key = wrt_key.tangent_of(pass); let tangent_id = builder.add_input(tangent_key); + debug_assert_eq!(tangent_id, value_kinds.len()); + value_kinds.push(LinearizedValueKind::Linear { + primal: Some(ValueKey::Input(wrt_key.clone())), + }); tangent_env.insert(ValueKey::Input(wrt_key.clone()), Some(tangent_id)); tangent_inputs.push((wrt_key.clone(), tangent_id)); } @@ -96,7 +103,8 @@ where continue; } - let mut primitive_builder = GraphPrimitiveBuilder::new(&mut builder); + let mut primitive_builder = + TrackingPrimitiveBuilder::new(&mut builder, &mut value_kinds); let tangent_out = operation.jvp_rule( &mut primitive_builder, &input_keys, @@ -116,8 +124,13 @@ where )); } - for (output_key, tangent_output) in output_keys.into_iter().zip(tangent_out) { - tangent_env.insert(output_key, tangent_output); + for (output_key, tangent_output) in output_keys.iter().zip(tangent_out) { + if let Some(tangent_id) = tangent_output { + value_kinds[tangent_id] = LinearizedValueKind::Linear { + primal: Some(output_key.clone()), + }; + } + tangent_env.insert(output_key.clone(), tangent_output); } } } @@ -136,9 +149,69 @@ where builder.build(), tangent_inputs, tangent_outputs, + value_kinds, )) } +struct TrackingPrimitiveBuilder<'a, Op: GraphOperation> { + inner: &'a mut GraphBuilder, + value_kinds: &'a mut Vec>, +} + +impl<'a, Op: GraphOperation> TrackingPrimitiveBuilder<'a, Op> { + fn new( + inner: &'a mut GraphBuilder, + value_kinds: &'a mut Vec>, + ) -> Self { + Self { inner, value_kinds } + } + + fn output_kind( + &self, + inputs: &[PrimitiveValue], + role: &OperationRole, + ) -> LinearizedValueKind { + let OperationRole::Linearized { active_mask } = role else { + return LinearizedValueKind::Residual; + }; + + if inputs.iter().enumerate().any(|(index, input)| { + active_mask.get(index).copied().unwrap_or(false) && self.is_linear_input(input) + }) { + LinearizedValueKind::Linear { primal: None } + } else { + LinearizedValueKind::Residual + } + } + + fn is_linear_input(&self, input: &PrimitiveValue) -> bool { + match input { + PrimitiveValue::Local(id) => { + matches!(self.value_kinds[*id], LinearizedValueKind::Linear { .. }) + } + PrimitiveValue::External(_) => true, + } + } +} + +impl PrimitiveBuilder for TrackingPrimitiveBuilder<'_, Op> { + fn add_primitive( + &mut self, + op: Op, + inputs: Vec>, + role: OperationRole, + ) -> Vec { + let output_kind = self.output_kind(&inputs, &role); + let inputs = inputs.into_iter().map(ValueRef::from).collect(); + let outputs = self.inner.add_operation(op, inputs, role); + for output_id in &outputs { + debug_assert_eq!(*output_id, self.value_kinds.len()); + self.value_kinds.push(output_kind.clone()); + } + outputs + } +} + fn output_keys( op_key: &OperationKey, output_count: usize, diff --git a/src/linearized_graph.rs b/src/linearized_graph.rs index dd44b0d..c24c76b 100644 --- a/src/linearized_graph.rs +++ b/src/linearized_graph.rs @@ -1,11 +1,24 @@ use computegraph::graph::Graph; -use computegraph::{GraphOperation, LocalValueId}; +use computegraph::{GraphOperation, LocalValueId, ValueKey}; + +/// Classification of a value inside a linearized graph. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum LinearizedValueKind { + /// Value independent of differentiated tangent flow. + Residual, + /// Value on differentiated tangent flow. + Linear { + /// Primal value whose tangent this value represents, when known. + primal: Option>, + }, +} /// Graph produced by linearizing a primitive computation graph. pub struct LinearizedGraph { graph: Graph, tangent_inputs: Vec<(Op::InputKey, LocalValueId)>, tangent_outputs: Vec>, + value_kinds: Vec>, } impl LinearizedGraph { @@ -13,11 +26,14 @@ impl LinearizedGraph { graph: Graph, tangent_inputs: Vec<(Op::InputKey, LocalValueId)>, tangent_outputs: Vec>, + value_kinds: Vec>, ) -> Self { + debug_assert_eq!(graph.values().len(), value_kinds.len()); Self { graph, tangent_inputs, tangent_outputs, + value_kinds, } } @@ -40,4 +56,9 @@ impl LinearizedGraph { pub fn tangent_outputs(&self) -> &[Option] { &self.tangent_outputs } + + /// Return the classification for a local value in the linearized graph. + pub fn value_kind(&self, id: LocalValueId) -> &LinearizedValueKind { + &self.value_kinds[id] + } } diff --git a/src/rules/mod.rs b/src/rules/mod.rs index 011f336..9a5c1b9 100644 --- a/src/rules/mod.rs +++ b/src/rules/mod.rs @@ -13,6 +13,5 @@ mod primitive_op; pub use ad_key::{ADKey, DiffPassId}; pub use ad_rule_error::{ADRuleError, ADRuleKind, ADRuleResult}; -pub(crate) use primitive_builder::GraphPrimitiveBuilder; -pub use primitive_builder::{PrimitiveBuilder, PrimitiveValue}; +pub use primitive_builder::{PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue}; pub use primitive_op::Primitive; diff --git a/src/rules/primitive_builder.rs b/src/rules/primitive_builder.rs index 4b7c3a4..bbdb7c7 100644 --- a/src/rules/primitive_builder.rs +++ b/src/rules/primitive_builder.rs @@ -1,5 +1,4 @@ -use computegraph::graph::GraphBuilder; -use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueRef}; +use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey, ValueRef}; /// Reference to a value available to a primitive AD rule. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -10,6 +9,52 @@ pub enum PrimitiveValue { External(computegraph::ValueKey), } +/// Input reference passed to a primitive transpose rule. +/// +/// Linearized graphs contain both residual values and values that belong to +/// the tangent flow. A transpose rule may freely use residual values as +/// ordinary graph operands, but tangent-flow values must not be retained as +/// tensor-valued operands in the transposed graph. When a tangent-flow value is +/// the tangent of a known primal value, `primal` carries that counterpart for +/// metadata or runtime-shape use by downstream primitive sets. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum PrimitiveTransposeInput { + /// Value that is independent of the differentiated tangent flow. + Residual(ValueKey), + /// Value on the differentiated tangent flow. + Linear { + /// Key of the linearized graph value. + key: ValueKey, + /// Primal value whose tangent this value represents, when known. + primal: Option>, + }, +} + +impl PrimitiveTransposeInput { + /// Return the key of the linearized graph value. + pub fn key(&self) -> &ValueKey { + match self { + Self::Residual(key) | Self::Linear { key, .. } => key, + } + } + + /// Return the corresponding primal key for a linear input, when known. + pub fn primal(&self) -> Option<&ValueKey> { + match self { + Self::Residual(_) => None, + Self::Linear { primal, .. } => primal.as_ref(), + } + } + + /// Return this input as a value usable by rule builders when it is residual. + pub fn as_residual_value(&self) -> Option> { + match self { + Self::Residual(key) => Some(PrimitiveValue::External(key.clone())), + Self::Linear { .. } => None, + } + } +} + impl From> for ValueRef { fn from(value: PrimitiveValue) -> Self { match value { @@ -38,25 +83,3 @@ pub trait PrimitiveBuilder { role: OperationRole, ) -> Vec; } - -pub(crate) struct GraphPrimitiveBuilder<'a, Op: GraphOperation> { - inner: &'a mut GraphBuilder, -} - -impl<'a, Op: GraphOperation> GraphPrimitiveBuilder<'a, Op> { - pub(crate) fn new(inner: &'a mut GraphBuilder) -> Self { - Self { inner } - } -} - -impl PrimitiveBuilder for GraphPrimitiveBuilder<'_, Op> { - fn add_primitive( - &mut self, - op: Op, - inputs: Vec>, - role: OperationRole, - ) -> Vec { - let inputs = inputs.into_iter().map(ValueRef::from).collect(); - self.inner.add_operation(op, inputs, role) - } -} diff --git a/src/rules/primitive_op.rs b/src/rules/primitive_op.rs index e44e6a4..8bfa77a 100644 --- a/src/rules/primitive_op.rs +++ b/src/rules/primitive_op.rs @@ -1,4 +1,4 @@ -use super::{ADKey, ADRuleResult, PrimitiveBuilder, PrimitiveValue}; +use super::{ADKey, ADRuleResult, PrimitiveBuilder, PrimitiveTransposeInput}; use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey}; /// Extends `GraphOperation` with primitive JVP and transpose rules for AD. @@ -14,7 +14,9 @@ use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey}; /// /// ``` /// use computegraph::{ValueKey, GraphOperation, LocalValueId, OperationRole}; -/// use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue}; +/// use tidu::{ +/// ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, +/// }; /// /// #[derive(Clone, Debug, PartialEq, Eq, Hash)] /// enum Key { Base(String), Tan(Box, DiffPassId) } @@ -48,7 +50,7 @@ use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey}; /// } /// fn transpose_rule( /// &self, _builder: &mut impl PrimitiveBuilder, -/// ct: &[Option], _i: &[PrimitiveValue], _m: &OperationRole, +/// ct: &[Option], _i: &[PrimitiveTransposeInput], _m: &OperationRole, /// _ctx: &mut (), /// ) -> tidu::ADRuleResult>> { /// Ok(vec![ct[0], ct[0]]) @@ -95,7 +97,7 @@ where &self, builder: &mut impl PrimitiveBuilder, cotangent_outputs: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, ctx: &mut Self::ADContext, ) -> ADRuleResult>> diff --git a/tests/adcontext_tests.rs b/tests/adcontext_tests.rs index 5bfa060..69f21bb 100644 --- a/tests/adcontext_tests.rs +++ b/tests/adcontext_tests.rs @@ -8,7 +8,9 @@ use computegraph::graph::{Graph, GraphBuilder}; use computegraph::resolve::resolve; use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef}; use computegraph::GraphOperation; -use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue}; +use tidu::{ + ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, +}; define_ad_key!(CtxKey); @@ -80,7 +82,7 @@ impl Primitive for CountingOp { &self, _builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - _inputs: &[PrimitiveValue], + _inputs: &[PrimitiveTransposeInput], _mode: &OperationRole, ctx: &mut CountingContext, ) -> tidu::ADRuleResult>> { diff --git a/tests/common/linearize_macros.rs b/tests/common/linearize_macros.rs index c84463c..0862d1e 100644 --- a/tests/common/linearize_macros.rs +++ b/tests/common/linearize_macros.rs @@ -1,7 +1,7 @@ #[allow(unused_imports)] use computegraph::types::{LocalValueId, OperationRole, ValueKey}; #[allow(unused_imports)] -use tidu::PrimitiveValue; +use tidu::{PrimitiveTransposeInput, PrimitiveValue}; #[macro_export] macro_rules! linearize_add { diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 7561bde..10eeade 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -8,7 +8,9 @@ use computegraph::resolve::{resolve, ResolvedView}; use computegraph::types::{LocalValueId, OperationRole, ValueKey}; use computegraph::{EvaluableGraphOperation, GraphOperation}; use tidu::LinearizedGraph; -use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue}; +use tidu::{ + ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, +}; use crate::{ define_ad_key, linearize_add, linearize_exp, linearize_mul, linearize_neg, transpose_add, @@ -90,7 +92,7 @@ impl Primitive for ScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { diff --git a/tests/common/transpose_macros.rs b/tests/common/transpose_macros.rs index 78b6baf..a5e62be 100644 --- a/tests/common/transpose_macros.rs +++ b/tests/common/transpose_macros.rs @@ -1,7 +1,7 @@ #[allow(unused_imports)] use computegraph::types::OperationRole; #[allow(unused_imports)] -use tidu::PrimitiveValue; +use tidu::{PrimitiveTransposeInput, PrimitiveValue}; #[macro_export] macro_rules! transpose_add { @@ -19,9 +19,10 @@ macro_rules! transpose_mul_real { }; let mut result = vec![None, None]; if active_mask[0] { + let fixed = $inputs[1].as_residual_value().unwrap(); let out = $builder.add_primitive( $OpMul, - vec![$inputs[1].clone(), PrimitiveValue::Local($ct)], + vec![fixed, PrimitiveValue::Local($ct)], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -29,9 +30,10 @@ macro_rules! transpose_mul_real { result[0] = Some(out[0]); } if active_mask[1] { + let fixed = $inputs[0].as_residual_value().unwrap(); let out = $builder.add_primitive( $OpMul, - vec![$inputs[0].clone(), PrimitiveValue::Local($ct)], + vec![fixed, PrimitiveValue::Local($ct)], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -51,9 +53,10 @@ macro_rules! transpose_mul_complex { }; let mut result = vec![None, None]; if active_mask[0] { + let fixed = $inputs[1].as_residual_value().unwrap(); let conj_fixed = $builder.add_primitive( $OpConj, - vec![$inputs[1].clone()], + vec![fixed], OperationRole::Linearized { active_mask: vec![false], }, @@ -71,9 +74,10 @@ macro_rules! transpose_mul_complex { result[0] = Some(out[0]); } if active_mask[1] { + let fixed = $inputs[0].as_residual_value().unwrap(); let conj_fixed = $builder.add_primitive( $OpConj, - vec![$inputs[0].clone()], + vec![fixed], OperationRole::Linearized { active_mask: vec![false], }, diff --git a/tests/complex_ad_tests.rs b/tests/complex_ad_tests.rs index 9a5af7b..721b943 100644 --- a/tests/complex_ad_tests.rs +++ b/tests/complex_ad_tests.rs @@ -10,7 +10,9 @@ use computegraph::resolve::resolve; use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef}; use computegraph::{EvaluableGraphOperation, GraphOperation}; use num_complex::Complex64; -use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue}; +use tidu::{ + ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, +}; const TOL: f64 = 1e-10; const NUM_TOL: f64 = 1e-5; @@ -115,7 +117,7 @@ impl Primitive for ComplexScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { diff --git a/tests/eager_backward_tests.rs b/tests/eager_backward_tests.rs index 47a428c..a233e05 100644 --- a/tests/eager_backward_tests.rs +++ b/tests/eager_backward_tests.rs @@ -11,7 +11,7 @@ use tidu::eager::{ }; use tidu::{ linear_transpose_with_builder, linearize, ADKey, DiffPassId, LinearizedGraph, Primitive, - PrimitiveBuilder, PrimitiveGraph, PrimitiveValue, + PrimitiveBuilder, PrimitiveGraph, PrimitiveTransposeInput, PrimitiveValue, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -141,7 +141,7 @@ impl Primitive for ScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -201,7 +201,7 @@ fn scalar_sum_terms( fn scalar_transpose_mul( builder: &mut impl PrimitiveBuilder, - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], ct: LocalValueId, role: &OperationRole, ) -> Vec> { @@ -211,9 +211,10 @@ fn scalar_transpose_mul( }; let mut result = vec![None, None]; if active_mask[0] { + let fixed = inputs[1].as_residual_value().unwrap(); let out = builder.add_primitive( ScalarOp::Mul, - vec![inputs[1].clone(), PrimitiveValue::Local(ct)], + vec![fixed, PrimitiveValue::Local(ct)], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -221,9 +222,10 @@ fn scalar_transpose_mul( result[0] = Some(out[0]); } if active_mask[1] { + let fixed = inputs[0].as_residual_value().unwrap(); let out = builder.add_primitive( ScalarOp::Mul, - vec![inputs[0].clone(), PrimitiveValue::Local(ct)], + vec![fixed, PrimitiveValue::Local(ct)], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -320,7 +322,7 @@ impl Primitive for TwoOutputOp { &self, _builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - _inputs: &[PrimitiveValue], + _inputs: &[PrimitiveTransposeInput], _mode: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { diff --git a/tests/eager_record_tests.rs b/tests/eager_record_tests.rs index 3efc366..dc65407 100644 --- a/tests/eager_record_tests.rs +++ b/tests/eager_record_tests.rs @@ -8,7 +8,7 @@ use tidu::eager::{ }; use tidu::{ linear_transpose_with_builder, ADKey, DiffPassId, LinearizedGraph, Primitive, PrimitiveBuilder, - PrimitiveGraph, PrimitiveValue, + PrimitiveGraph, PrimitiveTransposeInput, PrimitiveValue, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -150,7 +150,7 @@ impl Primitive for RecorderOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -212,7 +212,7 @@ fn sum_terms( fn transpose_mul( builder: &mut impl PrimitiveBuilder, - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], ct: LocalValueId, role: &OperationRole, ) -> Vec> { @@ -222,9 +222,10 @@ fn transpose_mul( }; let mut result = vec![None, None]; if active_mask[0] { + let fixed = inputs[1].as_residual_value().unwrap(); let out = builder.add_primitive( RecorderOp::Mul, - vec![inputs[1].clone(), PrimitiveValue::Local(ct)], + vec![fixed, PrimitiveValue::Local(ct)], OperationRole::Linearized { active_mask: vec![false, true], }, @@ -232,9 +233,10 @@ fn transpose_mul( result[0] = Some(out[0]); } if active_mask[1] { + let fixed = inputs[0].as_residual_value().unwrap(); let out = builder.add_primitive( RecorderOp::Mul, - vec![inputs[0].clone(), PrimitiveValue::Local(ct)], + vec![fixed, PrimitiveValue::Local(ct)], OperationRole::Linearized { active_mask: vec![false, true], }, diff --git a/tests/edge_case_tests.rs b/tests/edge_case_tests.rs index 1447373..4c75987 100644 --- a/tests/edge_case_tests.rs +++ b/tests/edge_case_tests.rs @@ -16,7 +16,9 @@ use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef}; use computegraph::{EvaluableGraphOperation, GraphOperation}; use ndarray::{ArrayD, Axis, IxDyn}; use num_complex::Complex64; -use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue}; +use tidu::{ + ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, +}; const TOL: f64 = 1e-10; const NUM_TOL: f64 = 1e-5; @@ -145,7 +147,7 @@ impl Primitive for ExtScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -357,7 +359,7 @@ impl Primitive for VectorOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -603,7 +605,7 @@ impl Primitive for ComplexVectorOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { diff --git a/tests/fallible_ad_tests.rs b/tests/fallible_ad_tests.rs index 4e15be6..a9a181c 100644 --- a/tests/fallible_ad_tests.rs +++ b/tests/fallible_ad_tests.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use computegraph::{GraphOperation, LocalValueId, OperationRole}; use tidu::{ linear_transpose, linear_transpose_with_builder, linearize, LinearizedGraph, PrimitiveBuilder, - PrimitiveValue, + PrimitiveTransposeInput, PrimitiveValue, }; use tidu::{ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive}; @@ -97,7 +97,7 @@ impl Primitive for Op { &self, _builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - _inputs: &[PrimitiveValue], + _inputs: &[PrimitiveTransposeInput], _mode: &OperationRole, _ctx: &mut (), ) -> ADRuleResult>> { diff --git a/tests/robustness_tests.rs b/tests/robustness_tests.rs index 1f35c23..5a4eb16 100644 --- a/tests/robustness_tests.rs +++ b/tests/robustness_tests.rs @@ -16,7 +16,9 @@ use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef}; use computegraph::{EvaluableGraphOperation, GraphOperation}; use ndarray::{ArrayD, IxDyn}; use num_complex::Complex64; -use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue}; +use tidu::{ + ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, +}; const TOL: f64 = 1e-10; @@ -202,7 +204,7 @@ impl Primitive for ComplexScalarOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { @@ -339,7 +341,7 @@ impl Primitive for VectorOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { diff --git a/tests/rules_public_api_tests.rs b/tests/rules_public_api_tests.rs index 7e84b2f..2769f75 100644 --- a/tests/rules_public_api_tests.rs +++ b/tests/rules_public_api_tests.rs @@ -7,7 +7,7 @@ use tidu::rules::{ }; use tidu::{ ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive, PrimitiveBuilder, - PrimitiveValue, + PrimitiveTransposeInput, PrimitiveValue, }; #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -64,7 +64,7 @@ impl Primitive for AddOp { &self, _builder: &mut impl PrimitiveBuilder, cotangent_outputs: &[Option], - _inputs: &[PrimitiveValue], + _inputs: &[PrimitiveTransposeInput], _mode: &OperationRole, _ctx: &mut Self::ADContext, ) -> tidu::ADRuleResult>> { diff --git a/tests/transpose_input_classification_tests.rs b/tests/transpose_input_classification_tests.rs new file mode 100644 index 0000000..cae0840 --- /dev/null +++ b/tests/transpose_input_classification_tests.rs @@ -0,0 +1,163 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use computegraph::graph::GraphBuilder; +use computegraph::resolve::resolve; +use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey}; +use tidu::{ + linear_transpose, linearize, ADKey, ADRuleResult, DiffPassId, Primitive, PrimitiveBuilder, + PrimitiveTransposeInput, PrimitiveValue, +}; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +enum Key { + User(u64), + Tangent { of: Box, pass: DiffPassId }, +} + +impl ADKey for Key { + fn tangent_of(&self, pass: DiffPassId) -> Self { + Self::Tangent { + of: Box::new(self.clone()), + pass, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +enum Op { + Add, + Probe, + ResidualHelper, + TangentHelper, +} + +impl GraphOperation for Op { + type Operand = f64; + type Context = (); + type InputKey = Key; + + fn input_count(&self) -> usize { + match self { + Self::Add => 2, + Self::Probe | Self::ResidualHelper | Self::TangentHelper => 1, + } + } + + fn output_count(&self) -> usize { + 1 + } +} + +#[derive(Default)] +struct SeenTransposeInputs { + add_inputs: Vec>>, +} + +impl Primitive for Op { + type ADContext = SeenTransposeInputs; + + fn add() -> Self { + Self::Add + } + + fn jvp_rule( + &self, + builder: &mut impl PrimitiveBuilder, + primal_in: &[ValueKey], + _primal_out: &[ValueKey], + tangent_in: &[Option], + _ctx: &mut Self::ADContext, + ) -> ADRuleResult>> { + match self { + Self::Probe => { + let Some(dx) = tangent_in[0] else { + return Ok(vec![None]); + }; + let residual = builder.add_primitive( + Self::ResidualHelper, + vec![PrimitiveValue::External(primal_in[0].clone())], + OperationRole::Linearized { + active_mask: vec![false], + }, + )[0]; + let tangent = builder.add_primitive( + Self::TangentHelper, + vec![PrimitiveValue::Local(dx)], + OperationRole::Linearized { + active_mask: vec![true], + }, + )[0]; + let mixed = builder.add_primitive( + Self::Add, + vec![ + PrimitiveValue::Local(residual), + PrimitiveValue::Local(tangent), + ], + OperationRole::Linearized { + active_mask: vec![false, true], + }, + )[0]; + Ok(vec![Some(mixed)]) + } + _ => Ok(vec![tangent_in[0]]), + } + } + + fn transpose_rule( + &self, + _builder: &mut impl PrimitiveBuilder, + cotangent_out: &[Option], + inputs: &[PrimitiveTransposeInput], + mode: &OperationRole, + ctx: &mut Self::ADContext, + ) -> ADRuleResult>> { + if matches!(self, Self::Add) { + ctx.add_inputs.push(inputs.to_vec()); + } + let Some(ct) = cotangent_out[0] else { + return Ok(vec![None; self.input_count()]); + }; + let active_mask = match mode { + OperationRole::Linearized { active_mask } => active_mask, + OperationRole::Primary => return Ok(vec![None; self.input_count()]), + }; + Ok(active_mask + .iter() + .map(|active| active.then_some(ct)) + .collect()) + } +} + +#[test] +fn linear_transpose_marks_residual_and_tangent_inputs_distinctly() { + let mut builder = GraphBuilder::::new(); + let x_key = Key::User(0); + let x = builder.add_input(x_key.clone()); + let y = builder.add_operation( + Op::Probe, + vec![computegraph::ValueRef::Local(x)], + OperationRole::Primary, + )[0]; + builder.set_outputs(vec![y]); + let graph = Arc::new(builder.build()); + let y_key = graph.values()[y].key.clone(); + let view = resolve(vec![graph]); + + let mut ctx = SeenTransposeInputs::default(); + let linear = linearize(&view, &[y_key], &[x_key], 0, &mut ctx, &HashMap::new()).unwrap(); + let _ = linear_transpose(&linear, &mut ctx).unwrap(); + + let add_inputs = ctx + .add_inputs + .first() + .expect("transposed Add rule should have been visited"); + assert!( + matches!(add_inputs[0], PrimitiveTransposeInput::Residual(_)), + "inactive residual helper must remain a residual transpose input: {add_inputs:?}" + ); + assert!( + matches!(add_inputs[1], PrimitiveTransposeInput::Linear { .. }), + "active tangent helper must be marked as a linear transpose input: {add_inputs:?}" + ); +} diff --git a/tests/vector_ad_tests.rs b/tests/vector_ad_tests.rs index 8e8242e..66f4b94 100644 --- a/tests/vector_ad_tests.rs +++ b/tests/vector_ad_tests.rs @@ -10,7 +10,9 @@ use computegraph::resolve::resolve; use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef}; use computegraph::{EvaluableGraphOperation, GraphOperation}; use ndarray::{ArrayD, Axis, IxDyn}; -use tidu::{ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue}; +use tidu::{ + ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, +}; const TOL: f64 = 1e-10; const NUM_TOL: f64 = 1e-5; @@ -190,7 +192,7 @@ impl Primitive for VectorOp { &self, builder: &mut impl PrimitiveBuilder, cotangent_out: &[Option], - inputs: &[PrimitiveValue], + inputs: &[PrimitiveTransposeInput], role: &OperationRole, _ctx: &mut (), ) -> tidu::ADRuleResult>> { From 57a2e7ebe7738ca2f8b5c96f4c6ce4e467b20495 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Tue, 7 Jul 2026 16:58:53 +0900 Subject: [PATCH 2/2] Split linearized graphs into residual and linear parts --- src/lib.rs | 3 +- src/linear_transpose.rs | 106 ++-------- src/linearize.rs | 118 +++-------- src/linearized_graph.rs | 54 ++--- src/split_builder.rs | 189 ++++++++++++++++++ tests/transpose_input_classification_tests.rs | 91 ++++++++- 6 files changed, 351 insertions(+), 210 deletions(-) create mode 100644 src/split_builder.rs diff --git a/src/lib.rs b/src/lib.rs index ba75f6e..514f507 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,10 +41,11 @@ mod linearize; mod linearized_graph; mod primitive_graph; pub mod rules; +mod split_builder; pub use linear_transpose::{linear_transpose, linear_transpose_with_builder}; pub use linearize::linearize; -pub use linearized_graph::{LinearizedGraph, LinearizedValueKind}; +pub use linearized_graph::LinearizedGraph; pub use primitive_graph::PrimitiveGraph; pub use rules::{ ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive, PrimitiveBuilder, diff --git a/src/linear_transpose.rs b/src/linear_transpose.rs index a4012c8..cab759c 100644 --- a/src/linear_transpose.rs +++ b/src/linear_transpose.rs @@ -1,12 +1,12 @@ use std::collections::HashMap; use crate::{ - ADKey, ADRuleError, ADRuleKind, ADRuleResult, LinearizedValueKind, Primitive, PrimitiveBuilder, + ADKey, ADRuleError, ADRuleKind, ADRuleResult, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, }; -use computegraph::graph::GraphBuilder; use computegraph::{LocalValueId, OperationRole, ValueKey, ValueRef}; +use crate::split_builder::SplitGraphBuilder; use crate::LinearizedGraph; /// Transpose a linearized graph, reversing linear flow. @@ -29,11 +29,10 @@ pub fn linear_transpose( where Op::InputKey: ADKey, { - let mut builder = GraphBuilder::::new(); + let mut builder = SplitGraphBuilder::::new(); let mut cotangent_env: HashMap, LocalValueId> = HashMap::new(); let mut cotangent_seed_inputs = Vec::new(); let graph = linear.as_graph(); - let mut value_kinds = Vec::new(); for (index, maybe_tangent_output) in linear.tangent_outputs().iter().enumerate() { let tangent_output_id = match maybe_tangent_output { @@ -43,11 +42,9 @@ where let source_key = graph.values()[*tangent_output_id].key.clone(); let seed_key = cotangent_seed_key(linear, index)?; - let seed_id = builder.add_input(seed_key.clone()); - debug_assert_eq!(seed_id, value_kinds.len()); - value_kinds.push(LinearizedValueKind::Linear { primal: None }); - cotangent_env.insert(source_key, seed_id); - cotangent_seed_inputs.push((seed_key, seed_id)); + let (seed_unified_id, seed_linear_id) = builder.add_linear_input(seed_key.clone(), None); + cotangent_env.insert(source_key, seed_unified_id); + cotangent_seed_inputs.push((seed_key, seed_linear_id)); } for op_node in graph.operations().iter().rev() { @@ -66,14 +63,16 @@ where .map(|input| transpose_input_for(linear, input)) .collect(); - let mut primitive_builder = TrackingPrimitiveBuilder::new(&mut builder, &mut value_kinds); let cotangent_in = op_node.operation.transpose_rule( - &mut primitive_builder, + &mut builder, &cotangent_out, &rule_inputs, &op_node.role, ctx, )?; + if let Some(err) = builder.take_error() { + return Err(err); + } if cotangent_in.len() != rule_inputs.len() { return Err(ADRuleError::invalid_input( format!("{:?}", op_node.operation), @@ -95,9 +94,7 @@ where match cotangent_env.get(&input_key).copied() { Some(existing_id) => { - let mut primitive_builder = - TrackingPrimitiveBuilder::new(&mut builder, &mut value_kinds); - let sum = primitive_builder.add_primitive( + let sum = builder.add_primitive( Op::add(), vec![ PrimitiveValue::Local(existing_id), @@ -121,19 +118,20 @@ where .iter() .map(|(_, tangent_input_id)| { let tangent_input_key = &graph.values()[*tangent_input_id].key; - cotangent_env.get(tangent_input_key).copied() + cotangent_env + .get(tangent_input_key) + .and_then(|unified_id| builder.linear_local_id(*unified_id)) }) .collect(); let active_outputs: Vec = tangent_outputs.iter().filter_map(|id| *id).collect(); - if !active_outputs.is_empty() { - builder.set_outputs(active_outputs); - } + let (linear_graph, residual_graph, linear_primals) = builder.finish(active_outputs); Ok(LinearizedGraph::from_parts( - builder.build(), + linear_graph, + residual_graph, cotangent_seed_inputs, tangent_outputs, - value_kinds, + linear_primals, )) } @@ -243,77 +241,15 @@ where match input { ValueRef::Local(local_id) => { let key = graph.values()[*local_id].key.clone(); - match linear.value_kind(*local_id) { - LinearizedValueKind::Residual => PrimitiveTransposeInput::Residual(key), - LinearizedValueKind::Linear { primal } => PrimitiveTransposeInput::Linear { - key, - primal: primal.clone(), - }, + PrimitiveTransposeInput::Linear { + key, + primal: linear.linear_primal(*local_id).cloned(), } } ValueRef::External(key) => PrimitiveTransposeInput::Residual(key.clone()), } } -struct TrackingPrimitiveBuilder<'a, Op: computegraph::GraphOperation> { - inner: &'a mut GraphBuilder, - value_kinds: &'a mut Vec>, -} - -impl<'a, Op: computegraph::GraphOperation> TrackingPrimitiveBuilder<'a, Op> { - fn new( - inner: &'a mut GraphBuilder, - value_kinds: &'a mut Vec>, - ) -> Self { - Self { inner, value_kinds } - } - - fn output_kind( - &self, - inputs: &[PrimitiveValue], - role: &OperationRole, - ) -> LinearizedValueKind { - let OperationRole::Linearized { active_mask } = role else { - return LinearizedValueKind::Residual; - }; - - if inputs.iter().enumerate().any(|(index, input)| { - active_mask.get(index).copied().unwrap_or(false) && self.is_linear_input(input) - }) { - LinearizedValueKind::Linear { primal: None } - } else { - LinearizedValueKind::Residual - } - } - - fn is_linear_input(&self, input: &PrimitiveValue) -> bool { - match input { - PrimitiveValue::Local(id) => { - matches!(self.value_kinds[*id], LinearizedValueKind::Linear { .. }) - } - PrimitiveValue::External(_) => true, - } - } -} - -impl PrimitiveBuilder for TrackingPrimitiveBuilder<'_, Op> { - fn add_primitive( - &mut self, - op: Op, - inputs: Vec>, - role: OperationRole, - ) -> Vec { - let output_kind = self.output_kind(&inputs, &role); - let inputs = inputs.into_iter().map(ValueRef::from).collect(); - let outputs = self.inner.add_operation(op, inputs, role); - for output_id in &outputs { - debug_assert_eq!(*output_id, self.value_kinds.len()); - self.value_kinds.push(output_kind.clone()); - } - outputs - } -} - fn cotangent_seed_key( linear: &LinearizedGraph, index: usize, diff --git a/src/linearize.rs b/src/linearize.rs index 0d0b5e5..cca1689 100644 --- a/src/linearize.rs +++ b/src/linearize.rs @@ -1,14 +1,11 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use crate::rules::{PrimitiveBuilder, PrimitiveValue}; -use crate::{ - ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, LinearizedValueKind, Primitive, -}; -use computegraph::graph::GraphBuilder; +use crate::{ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive}; use computegraph::resolve::{ResolvedView, ValueDef}; -use computegraph::{GraphOperation, LocalValueId, OperationKey, OperationRole, ValueKey, ValueRef}; +use computegraph::{GraphOperation, LocalValueId, OperationKey, ValueKey}; +use crate::split_builder::SplitGraphBuilder; use crate::LinearizedGraph; /// Linearize a resolved computation graph, producing a linear graph. @@ -41,22 +38,18 @@ pub fn linearize( where Op::InputKey: ADKey, { - let mut builder = GraphBuilder::::new(); + let mut builder = SplitGraphBuilder::::new(); let topo_keys = topological_order(view, outputs, aliases); let mut tangent_env: HashMap, Option> = HashMap::new(); let mut processed_ops = HashSet::new(); - let mut value_kinds = Vec::new(); let mut tangent_inputs = Vec::with_capacity(wrt.len()); for wrt_key in wrt { let tangent_key = wrt_key.tangent_of(pass); - let tangent_id = builder.add_input(tangent_key); - debug_assert_eq!(tangent_id, value_kinds.len()); - value_kinds.push(LinearizedValueKind::Linear { - primal: Some(ValueKey::Input(wrt_key.clone())), - }); - tangent_env.insert(ValueKey::Input(wrt_key.clone()), Some(tangent_id)); - tangent_inputs.push((wrt_key.clone(), tangent_id)); + let (unified_id, linear_id) = + builder.add_linear_input(tangent_key, Some(ValueKey::Input(wrt_key.clone()))); + tangent_env.insert(ValueKey::Input(wrt_key.clone()), Some(unified_id)); + tangent_inputs.push((wrt_key.clone(), linear_id)); } for key in topo_keys { @@ -103,15 +96,16 @@ where continue; } - let mut primitive_builder = - TrackingPrimitiveBuilder::new(&mut builder, &mut value_kinds); let tangent_out = operation.jvp_rule( - &mut primitive_builder, + &mut builder, &input_keys, &output_keys, &tangent_in, ctx, )?; + if let Some(err) = builder.take_error() { + return Err(err); + } if tangent_out.len() != output_keys.len() { return Err(ADRuleError::invalid_input( format!("{:?}", operation), @@ -126,9 +120,7 @@ where for (output_key, tangent_output) in output_keys.iter().zip(tangent_out) { if let Some(tangent_id) = tangent_output { - value_kinds[tangent_id] = LinearizedValueKind::Linear { - primal: Some(output_key.clone()), - }; + builder.set_linear_primal(tangent_id, output_key.clone())?; } tangent_env.insert(output_key.clone(), tangent_output); } @@ -138,80 +130,32 @@ where let tangent_outputs: Vec> = outputs .iter() - .map(|key| tangent_env.get(key).copied().flatten()) - .collect(); + .map(|key| match tangent_env.get(key).copied().flatten() { + Some(unified_id) => builder + .linear_local_id(unified_id) + .ok_or_else(|| { + ADRuleError::invalid_input( + "tidu::linearize", + ADRuleKind::Jvp, + "JVP rule returned a residual value as an active requested tangent output", + ) + }) + .map(Some), + None => Ok(None), + }) + .collect::>>()?; let active_outputs: Vec = tangent_outputs.iter().filter_map(|id| *id).collect(); - if !active_outputs.is_empty() { - builder.set_outputs(active_outputs); - } + let (linear, residual, linear_primals) = builder.finish(active_outputs); Ok(LinearizedGraph::from_parts( - builder.build(), + linear, + residual, tangent_inputs, tangent_outputs, - value_kinds, + linear_primals, )) } -struct TrackingPrimitiveBuilder<'a, Op: GraphOperation> { - inner: &'a mut GraphBuilder, - value_kinds: &'a mut Vec>, -} - -impl<'a, Op: GraphOperation> TrackingPrimitiveBuilder<'a, Op> { - fn new( - inner: &'a mut GraphBuilder, - value_kinds: &'a mut Vec>, - ) -> Self { - Self { inner, value_kinds } - } - - fn output_kind( - &self, - inputs: &[PrimitiveValue], - role: &OperationRole, - ) -> LinearizedValueKind { - let OperationRole::Linearized { active_mask } = role else { - return LinearizedValueKind::Residual; - }; - - if inputs.iter().enumerate().any(|(index, input)| { - active_mask.get(index).copied().unwrap_or(false) && self.is_linear_input(input) - }) { - LinearizedValueKind::Linear { primal: None } - } else { - LinearizedValueKind::Residual - } - } - - fn is_linear_input(&self, input: &PrimitiveValue) -> bool { - match input { - PrimitiveValue::Local(id) => { - matches!(self.value_kinds[*id], LinearizedValueKind::Linear { .. }) - } - PrimitiveValue::External(_) => true, - } - } -} - -impl PrimitiveBuilder for TrackingPrimitiveBuilder<'_, Op> { - fn add_primitive( - &mut self, - op: Op, - inputs: Vec>, - role: OperationRole, - ) -> Vec { - let output_kind = self.output_kind(&inputs, &role); - let inputs = inputs.into_iter().map(ValueRef::from).collect(); - let outputs = self.inner.add_operation(op, inputs, role); - for output_id in &outputs { - debug_assert_eq!(*output_id, self.value_kinds.len()); - self.value_kinds.push(output_kind.clone()); - } - outputs - } -} - fn output_keys( op_key: &OperationKey, output_count: usize, diff --git a/src/linearized_graph.rs b/src/linearized_graph.rs index c24c76b..c838ae5 100644 --- a/src/linearized_graph.rs +++ b/src/linearized_graph.rs @@ -1,50 +1,53 @@ +use std::sync::Arc; + use computegraph::graph::Graph; use computegraph::{GraphOperation, LocalValueId, ValueKey}; -/// Classification of a value inside a linearized graph. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub enum LinearizedValueKind { - /// Value independent of differentiated tangent flow. - Residual, - /// Value on differentiated tangent flow. - Linear { - /// Primal value whose tangent this value represents, when known. - primal: Option>, - }, -} - /// Graph produced by linearizing a primitive computation graph. pub struct LinearizedGraph { - graph: Graph, + linear: Graph, + residual: Arc>, tangent_inputs: Vec<(Op::InputKey, LocalValueId)>, tangent_outputs: Vec>, - value_kinds: Vec>, + linear_primals: Vec>>, } impl LinearizedGraph { pub(crate) fn from_parts( - graph: Graph, + linear: Graph, + residual: Arc>, tangent_inputs: Vec<(Op::InputKey, LocalValueId)>, tangent_outputs: Vec>, - value_kinds: Vec>, + linear_primals: Vec>>, ) -> Self { - debug_assert_eq!(graph.values().len(), value_kinds.len()); + debug_assert_eq!(linear.values().len(), linear_primals.len()); Self { - graph, + linear, + residual, tangent_inputs, tangent_outputs, - value_kinds, + linear_primals, } } - /// Borrow the lower-level graph representation. + /// Borrow the strictly-linear graph representation. pub fn as_graph(&self) -> &Graph { - &self.graph + &self.linear } - /// Consume this value and return the lower-level graph representation. + /// Borrow the residual graph referenced by the linear graph. + pub fn residual_graph(&self) -> &Graph { + &self.residual + } + + /// Consume this value and return the strictly-linear graph representation. pub fn into_graph(self) -> Graph { - self.graph + self.linear + } + + /// Consume this value and return both graph parts. + pub fn into_graphs(self) -> (Graph, Arc>) { + (self.linear, self.residual) } /// Tangent input keys and local value ids. @@ -57,8 +60,7 @@ impl LinearizedGraph { &self.tangent_outputs } - /// Return the classification for a local value in the linearized graph. - pub fn value_kind(&self, id: LocalValueId) -> &LinearizedValueKind { - &self.value_kinds[id] + pub(crate) fn linear_primal(&self, id: LocalValueId) -> Option<&ValueKey> { + self.linear_primals.get(id).and_then(Option::as_ref) } } diff --git a/src/split_builder.rs b/src/split_builder.rs new file mode 100644 index 0000000..c5f25d9 --- /dev/null +++ b/src/split_builder.rs @@ -0,0 +1,189 @@ +use std::sync::Arc; + +use crate::{ADRuleError, ADRuleKind, PrimitiveBuilder, PrimitiveValue}; +use computegraph::graph::{Graph, GraphBuilder}; +use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey, ValueRef}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum SplitValue { + Residual(LocalValueId), + Linear(LocalValueId), +} + +pub(crate) struct SplitGraphBuilder { + residual: GraphBuilder, + linear: GraphBuilder, + unified_values: Vec, + linear_primals: Vec>>, + first_error: Option, +} + +impl SplitGraphBuilder { + pub(crate) fn new() -> Self { + Self { + residual: GraphBuilder::new(), + linear: GraphBuilder::new(), + unified_values: Vec::new(), + linear_primals: Vec::new(), + first_error: None, + } + } + + pub(crate) fn add_linear_input( + &mut self, + key: Op::InputKey, + primal: Option>, + ) -> (LocalValueId, LocalValueId) { + let local_id = self.linear.add_input(key); + self.push_linear_value(local_id, primal); + let unified_id = self.unified_values.len() - 1; + (unified_id, local_id) + } + + pub(crate) fn linear_local_id(&self, unified_id: LocalValueId) -> Option { + match self.unified_values.get(unified_id).copied() { + Some(SplitValue::Linear(local_id)) => Some(local_id), + _ => None, + } + } + + pub(crate) fn set_linear_primal( + &mut self, + unified_id: LocalValueId, + primal: ValueKey, + ) -> Result<(), ADRuleError> { + let Some(SplitValue::Linear(local_id)) = self.unified_values.get(unified_id).copied() + else { + return Err(ADRuleError::invalid_input( + "tidu::linearize", + ADRuleKind::Jvp, + "JVP rule returned a residual value as an active tangent output", + )); + }; + if let Some(slot) = self.linear_primals.get_mut(local_id) { + *slot = Some(primal); + } + Ok(()) + } + + pub(crate) fn take_error(&mut self) -> Option { + self.first_error.take() + } + + pub(crate) fn finish( + mut self, + linear_outputs: Vec, + ) -> (Graph, Arc>, Vec>>) { + if !linear_outputs.is_empty() { + self.linear.set_outputs(linear_outputs); + } + let residual = Arc::new(self.residual.build()); + self.linear.add_parent(Arc::clone(&residual)); + (self.linear.build(), residual, self.linear_primals) + } + + fn push_residual_value(&mut self, local_id: LocalValueId) { + self.unified_values.push(SplitValue::Residual(local_id)); + } + + fn push_linear_value(&mut self, local_id: LocalValueId, primal: Option>) { + debug_assert_eq!(local_id, self.linear_primals.len()); + self.unified_values.push(SplitValue::Linear(local_id)); + self.linear_primals.push(primal); + } + + fn is_linear_value(&self, input: &PrimitiveValue) -> bool { + match input { + PrimitiveValue::Local(local_id) => { + matches!( + self.unified_values.get(*local_id), + Some(SplitValue::Linear(_)) + ) + } + PrimitiveValue::External(_) => false, + } + } + + fn record_inactive_tangent_error(&mut self, op: &Op, input_index: usize) { + if self.first_error.is_none() { + self.first_error = Some(ADRuleError::invalid_input( + format!("{op:?}"), + ADRuleKind::Jvp, + format!( + "linearized operation used tangent-flow input {input_index} in an inactive position" + ), + )); + } + } + + fn validate_linearity(&mut self, op: &Op, inputs: &[PrimitiveValue], role: &OperationRole) { + let OperationRole::Linearized { active_mask } = role else { + return; + }; + for (input_index, input) in inputs.iter().enumerate() { + if self.is_linear_value(input) + && !active_mask.get(input_index).copied().unwrap_or(false) + { + self.record_inactive_tangent_error(op, input_index); + } + } + } + + fn input_for_residual_graph(&self, input: PrimitiveValue) -> ValueRef { + match input { + PrimitiveValue::Local(local_id) => match self.unified_values[local_id] { + SplitValue::Residual(real_id) => ValueRef::Local(real_id), + SplitValue::Linear(real_id) => { + ValueRef::External(self.linear.global_key(real_id).clone()) + } + }, + PrimitiveValue::External(key) => ValueRef::External(key), + } + } + + fn input_for_linear_graph(&self, input: PrimitiveValue) -> ValueRef { + match input { + PrimitiveValue::Local(local_id) => match self.unified_values[local_id] { + SplitValue::Residual(real_id) => { + ValueRef::External(self.residual.global_key(real_id).clone()) + } + SplitValue::Linear(real_id) => ValueRef::Local(real_id), + }, + PrimitiveValue::External(key) => ValueRef::External(key), + } + } +} + +impl PrimitiveBuilder for SplitGraphBuilder { + fn add_primitive( + &mut self, + op: Op, + inputs: Vec>, + role: OperationRole, + ) -> Vec { + self.validate_linearity(&op, &inputs, &role); + let target_linear = inputs.iter().any(|input| self.is_linear_value(input)); + if target_linear { + let graph_inputs = inputs + .into_iter() + .map(|input| self.input_for_linear_graph(input)) + .collect(); + let outputs = self.linear.add_operation(op, graph_inputs, role); + for &output_id in &outputs { + self.push_linear_value(output_id, None); + } + return (self.unified_values.len() - outputs.len()..self.unified_values.len()) + .collect(); + } + + let graph_inputs = inputs + .into_iter() + .map(|input| self.input_for_residual_graph(input)) + .collect(); + let outputs = self.residual.add_operation(op, graph_inputs, role); + for &output_id in &outputs { + self.push_residual_value(output_id); + } + (self.unified_values.len() - outputs.len()..self.unified_values.len()).collect() + } +} diff --git a/tests/transpose_input_classification_tests.rs b/tests/transpose_input_classification_tests.rs index cae0840..6c0ecf5 100644 --- a/tests/transpose_input_classification_tests.rs +++ b/tests/transpose_input_classification_tests.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use computegraph::graph::GraphBuilder; use computegraph::resolve::resolve; -use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey}; +use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey, ValueRef}; use tidu::{ linear_transpose, linearize, ADKey, ADRuleResult, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveTransposeInput, PrimitiveValue, @@ -28,6 +28,7 @@ impl ADKey for Key { enum Op { Add, Probe, + BadInactive, ResidualHelper, TangentHelper, } @@ -40,7 +41,7 @@ impl GraphOperation for Op { fn input_count(&self) -> usize { match self { Self::Add => 2, - Self::Probe | Self::ResidualHelper | Self::TangentHelper => 1, + Self::Probe | Self::BadInactive | Self::ResidualHelper | Self::TangentHelper => 1, } } @@ -70,10 +71,20 @@ impl Primitive for Op { _ctx: &mut Self::ADContext, ) -> ADRuleResult>> { match self { - Self::Probe => { + Self::Probe | Self::BadInactive => { let Some(dx) = tangent_in[0] else { return Ok(vec![None]); }; + if matches!(self, Self::BadInactive) { + let leaked = builder.add_primitive( + Self::TangentHelper, + vec![PrimitiveValue::Local(dx)], + OperationRole::Linearized { + active_mask: vec![false], + }, + )[0]; + return Ok(vec![Some(leaked)]); + } let residual = builder.add_primitive( Self::ResidualHelper, vec![PrimitiveValue::External(primal_in[0].clone())], @@ -129,21 +140,79 @@ impl Primitive for Op { } } -#[test] -fn linear_transpose_marks_residual_and_tangent_inputs_distinctly() { +fn resolved_probe_graph(op: Op) -> (computegraph::resolve::ResolvedView, ValueKey, Key) { let mut builder = GraphBuilder::::new(); let x_key = Key::User(0); let x = builder.add_input(x_key.clone()); - let y = builder.add_operation( - Op::Probe, - vec![computegraph::ValueRef::Local(x)], - OperationRole::Primary, - )[0]; + let y = builder.add_operation(op, vec![ValueRef::Local(x)], OperationRole::Primary)[0]; builder.set_outputs(vec![y]); let graph = Arc::new(builder.build()); let y_key = graph.values()[y].key.clone(); - let view = resolve(vec![graph]); + (resolve(vec![graph]), y_key, x_key) +} + +#[test] +fn linearize_splits_residual_helpers_out_of_the_linear_graph() { + let (view, y_key, x_key) = resolved_probe_graph(Op::Probe); + let mut ctx = SeenTransposeInputs::default(); + + let linear = linearize(&view, &[y_key], &[x_key], 0, &mut ctx, &HashMap::new()).unwrap(); + assert!( + linear + .residual_graph() + .operations() + .iter() + .any(|node| node.operation == Op::ResidualHelper), + "residual-only helper should be emitted to the residual graph" + ); + assert!( + !linear + .as_graph() + .operations() + .iter() + .any(|node| node.operation == Op::ResidualHelper), + "linear graph must not retain residual-only helpers" + ); + + let add = linear + .as_graph() + .operations() + .iter() + .find(|node| node.operation == Op::Add) + .expect("linear graph should contain the tangent/residual join"); + assert!( + matches!(add.inputs[0], ValueRef::External(_)), + "residual helper should be referenced by key from the linear graph: {:?}", + add.inputs + ); + assert!( + matches!(add.inputs[1], ValueRef::Local(_)), + "tangent helper should remain local to the linear graph: {:?}", + add.inputs + ); +} + +#[test] +fn linearize_rejects_tangent_flow_in_inactive_positions() { + let (view, y_key, x_key) = resolved_probe_graph(Op::BadInactive); + let mut ctx = SeenTransposeInputs::default(); + + let err = match linearize(&view, &[y_key], &[x_key], 0, &mut ctx, &HashMap::new()) { + Ok(_) => panic!("linearize should reject inactive tangent flow"), + Err(err) => err, + }; + + assert_eq!(err.rule(), tidu::ADRuleKind::Jvp); + assert!( + err.to_string().contains("inactive"), + "linearity violation should mention the inactive tangent operand: {err:?}" + ); +} + +#[test] +fn linear_transpose_marks_residual_and_tangent_inputs_distinctly() { + let (view, y_key, x_key) = resolved_probe_graph(Op::Probe); let mut ctx = SeenTransposeInputs::default(); let linear = linearize(&view, &[y_key], &[x_key], 0, &mut ctx, &HashMap::new()).unwrap(); let _ = linear_transpose(&linear, &mut ctx).unwrap();