[codex] classify linear transpose inputs

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<Self>,
         cotangent_outputs: &[Option<LocalValueId>],
-        inputs: &[PrimitiveValue<Self>],
+        inputs: &[PrimitiveTransposeInput<Self>],
         role: &OperationRole,
         _ctx: &mut (),
     ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
@@ -390,7 +390,7 @@ fn sum_tangent_terms(
 
 fn transpose_mul(
     builder: &mut impl PrimitiveBuilder<ScalarOp>,
-    inputs: &[PrimitiveValue<ScalarOp>],
+    inputs: &[PrimitiveTransposeInput<ScalarOp>],
     ct: LocalValueId,
     role: &OperationRole,
 ) -> Vec<Option<LocalValueId>> {
@@ -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],
             },

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<Self>,
         cotangent_outputs: &[Option<LocalValueId>],
-        inputs: &[PrimitiveValue<Self>],
+        inputs: &[PrimitiveTransposeInput<Self>],
         role: &OperationRole,
         _ctx: &mut (),
     ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
@@ -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],
                         },

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<Self>,
         cotangent_outputs: &[Option<LocalValueId>],
-        inputs: &[PrimitiveValue<Self>],
+        inputs: &[PrimitiveTransposeInput<Self>],
         role: &OperationRole,
         _ctx: &mut (),
     ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
@@ -216,7 +216,7 @@ fn sum_tangent_terms(
 
 fn transpose_mul(
     builder: &mut impl PrimitiveBuilder<ScalarOp>,
-    inputs: &[PrimitiveValue<ScalarOp>],
+    inputs: &[PrimitiveTransposeInput<ScalarOp>],
     ct: LocalValueId,
     role: &OperationRole,
 ) -> Vec<Option<LocalValueId>> {
@@ -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],
             },

src/eager/record.rs

@@ -404,7 +404,7 @@ impl<K> Recorder<K> {
     /// 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<K> Recorder<K> {
     ///         &self,
     ///         _builder: &mut impl PrimitiveBuilder<Self>,
     ///         cotangent_out: &[Option<LocalValueId>],
-    ///         _inputs: &[PrimitiveValue<Self>],
+    ///         _inputs: &[PrimitiveTransposeInput<Self>],
     ///         _role: &OperationRole,
     ///         _ctx: &mut (),
     ///     ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
@@ -582,7 +582,7 @@ fn fresh_value_keys<Op: GraphOperation>(
 #[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<Self>,
             cotangent_out: &[Option<LocalValueId>],
-            _inputs: &[PrimitiveValue<Self>],
+            _inputs: &[PrimitiveTransposeInput<Self>],
             _role: &OperationRole,
             _ctx: &mut (),
         ) -> ADRuleResult<Vec<Option<LocalValueId>>> {

src/lib.rs

@@ -41,12 +41,13 @@ 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;
 pub use primitive_graph::PrimitiveGraph;
 pub use rules::{
     ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive, PrimitiveBuilder,
-    PrimitiveValue,
+    PrimitiveTransposeInput, PrimitiveValue,
 };

src/linear_transpose.rs

@@ -1,12 +1,12 @@
 use std::collections::HashMap;
 
-use crate::rules::GraphPrimitiveBuilder;
 use crate::{
-    ADKey, ADRuleError, ADRuleKind, ADRuleResult, Primitive, PrimitiveBuilder, PrimitiveValue,
+    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,7 +29,7 @@ pub fn linear_transpose<Op: Primitive>(
 where
     Op::InputKey: ADKey,
 {
-    let mut builder = GraphBuilder::<Op>::new();
+    let mut builder = SplitGraphBuilder::<Op>::new();
     let mut cotangent_env: HashMap<ValueKey<Op>, LocalValueId> = HashMap::new();
     let mut cotangent_seed_inputs = Vec::new();
     let graph = linear.as_graph();
@@ -42,9 +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());
-        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() {
@@ -57,25 +57,22 @@ where
             continue;
         }
 
-        let rule_inputs: Vec<PrimitiveValue<Op>> = op_node
+        let rule_inputs: Vec<PrimitiveTransposeInput<Op>> = 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 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),
@@ -93,17 +90,11 @@ 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 sum = primitive_builder.add_primitive(
+                    let sum = builder.add_primitive(
                         Op::add(),
                         vec![
                             PrimitiveValue::Local(existing_id),
@@ -127,18 +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<LocalValueId> = 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,
+        linear_primals,
     ))
 }
 
@@ -174,15 +167,10 @@ where
             continue;
         }
 
-        let rule_inputs: Vec<PrimitiveValue<Op>> = op_node
+        let rule_inputs: Vec<PrimitiveTransposeInput<Op>> = 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 +197,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 +230,26 @@ where
         .collect())
 }
 
+fn transpose_input_for<Op: Primitive>(
+    linear: &LinearizedGraph<Op>,
+    input: &ValueRef<Op>,
+) -> PrimitiveTransposeInput<Op>
+where
+    Op::InputKey: ADKey,
+{
+    let graph = linear.as_graph();
+    match input {
+        ValueRef::Local(local_id) => {
+            let key = graph.values()[*local_id].key.clone();
+            PrimitiveTransposeInput::Linear {
+                key,
+                primal: linear.linear_primal(*local_id).cloned(),
+            }
+        }
+        ValueRef::External(key) => PrimitiveTransposeInput::Residual(key.clone()),
+    }
+}
+
 fn cotangent_seed_key<Op: Primitive>(
     linear: &LinearizedGraph<Op>,
     index: usize,