tidu/
linear_transpose.rs

1use std::collections::HashMap;
2
3use crate::{
4    ADKey, ADRuleError, ADRuleKind, ADRuleResult, Primitive, PrimitiveBuilder,
5    PrimitiveTransposeInput, PrimitiveValue,
6};
7use computegraph::{LocalValueId, OperationRole, ValueKey, ValueRef};
8
9use crate::split_builder::SplitGraphBuilder;
10use crate::LinearizedGraph;
11
12/// Transpose a linearized graph, reversing linear flow.
13///
14/// Fan-out accumulation is emitted explicitly with [`crate::Primitive::add`];
15/// no duplication primitive is assumed by the graph transform.
16///
17/// # Examples
18///
19/// ```ignore
20/// let mut ctx = ();
21/// let transposed = tidu::linear_transpose(&linear, &mut ctx)?;
22/// assert_eq!(transposed.tangent_outputs().len(), linear.tangent_inputs().len());
23/// # Ok::<(), tidu::ADRuleError>(())
24/// ```
25pub fn linear_transpose<Op: Primitive>(
26    linear: &LinearizedGraph<Op>,
27    ctx: &mut Op::ADContext,
28) -> ADRuleResult<LinearizedGraph<Op>>
29where
30    Op::InputKey: ADKey,
31{
32    let mut builder = SplitGraphBuilder::<Op>::new();
33    let mut cotangent_env: HashMap<ValueKey<Op>, LocalValueId> = HashMap::new();
34    let mut cotangent_seed_inputs = Vec::new();
35    let graph = linear.as_graph();
36
37    for (index, maybe_tangent_output) in linear.tangent_outputs().iter().enumerate() {
38        let tangent_output_id = match maybe_tangent_output {
39            Some(tangent_output_id) => tangent_output_id,
40            None => continue,
41        };
42
43        let source_key = graph.values()[*tangent_output_id].key.clone();
44        let seed_key = cotangent_seed_key(linear, index)?;
45        let (seed_unified_id, seed_linear_id) = builder.add_linear_input(seed_key.clone(), None);
46        cotangent_env.insert(source_key, seed_unified_id);
47        cotangent_seed_inputs.push((seed_key, seed_linear_id));
48    }
49
50    for op_node in graph.operations().iter().rev() {
51        let cotangent_out: Vec<Option<LocalValueId>> = op_node
52            .outputs
53            .iter()
54            .map(|output_id| cotangent_env.get(&graph.values()[*output_id].key).copied())
55            .collect();
56        if cotangent_out.iter().all(Option::is_none) {
57            continue;
58        }
59
60        let rule_inputs: Vec<PrimitiveTransposeInput<Op>> = op_node
61            .inputs
62            .iter()
63            .map(|input| transpose_input_for(linear, input))
64            .collect();
65
66        let cotangent_in = op_node.operation.transpose_rule(
67            &mut builder,
68            &cotangent_out,
69            &rule_inputs,
70            &op_node.role,
71            ctx,
72        )?;
73        if let Some(err) = builder.take_error() {
74            return Err(err);
75        }
76        if cotangent_in.len() != rule_inputs.len() {
77            return Err(ADRuleError::invalid_input(
78                format!("{:?}", op_node.operation),
79                ADRuleKind::Transpose,
80                format!(
81                    "rule returned {} cotangents for {} inputs",
82                    cotangent_in.len(),
83                    rule_inputs.len()
84                ),
85            ));
86        }
87
88        for (input, maybe_cotangent) in rule_inputs.iter().zip(cotangent_in) {
89            let cotangent_id = match maybe_cotangent {
90                Some(cotangent_id) => cotangent_id,
91                None => continue,
92            };
93            let input_key = input.key().clone();
94
95            match cotangent_env.get(&input_key).copied() {
96                Some(existing_id) => {
97                    let sum = builder.add_primitive(
98                        Op::add(),
99                        vec![
100                            PrimitiveValue::Local(existing_id),
101                            PrimitiveValue::Local(cotangent_id),
102                        ],
103                        OperationRole::Linearized {
104                            active_mask: vec![true, true],
105                        },
106                    );
107                    cotangent_env.insert(input_key, sum[0]);
108                }
109                None => {
110                    cotangent_env.insert(input_key, cotangent_id);
111                }
112            }
113        }
114    }
115
116    let tangent_outputs: Vec<Option<LocalValueId>> = linear
117        .tangent_inputs()
118        .iter()
119        .map(|(_, tangent_input_id)| {
120            let tangent_input_key = &graph.values()[*tangent_input_id].key;
121            cotangent_env
122                .get(tangent_input_key)
123                .and_then(|unified_id| builder.linear_local_id(*unified_id))
124        })
125        .collect();
126    let active_outputs: Vec<LocalValueId> = tangent_outputs.iter().filter_map(|id| *id).collect();
127    let (linear_graph, residual_graph, linear_primals) = builder.finish(active_outputs);
128
129    Ok(LinearizedGraph::from_parts(
130        linear_graph,
131        residual_graph,
132        cotangent_seed_inputs,
133        tangent_outputs,
134        linear_primals,
135    ))
136}
137
138/// Execute the transpose of a linearized graph using a caller-provided builder.
139pub fn linear_transpose_with_builder<Op: Primitive>(
140    linear: &LinearizedGraph<Op>,
141    builder: &mut impl PrimitiveBuilder<Op>,
142    cotangent_seeds: &[Option<LocalValueId>],
143    ctx: &mut Op::ADContext,
144) -> ADRuleResult<Vec<Option<LocalValueId>>>
145where
146    Op::InputKey: ADKey,
147{
148    let mut cotangent_env: HashMap<ValueKey<Op>, LocalValueId> = HashMap::new();
149    let graph = linear.as_graph();
150
151    for (index, maybe_tangent_output) in linear.tangent_outputs().iter().enumerate() {
152        if let (Some(output_id), Some(Some(seed_id))) =
153            (maybe_tangent_output, cotangent_seeds.get(index))
154        {
155            let key = graph.values()[*output_id].key.clone();
156            cotangent_env.insert(key, *seed_id);
157        }
158    }
159
160    for op_node in graph.operations().iter().rev() {
161        let cotangent_out: Vec<Option<LocalValueId>> = op_node
162            .outputs
163            .iter()
164            .map(|output_id| cotangent_env.get(&graph.values()[*output_id].key).copied())
165            .collect();
166        if cotangent_out.iter().all(Option::is_none) {
167            continue;
168        }
169
170        let rule_inputs: Vec<PrimitiveTransposeInput<Op>> = op_node
171            .inputs
172            .iter()
173            .map(|input| transpose_input_for(linear, input))
174            .collect();
175
176        let cotangent_in = op_node.operation.transpose_rule(
177            builder,
178            &cotangent_out,
179            &rule_inputs,
180            &op_node.role,
181            ctx,
182        )?;
183        if cotangent_in.len() != rule_inputs.len() {
184            return Err(ADRuleError::invalid_input(
185                format!("{:?}", op_node.operation),
186                ADRuleKind::Transpose,
187                format!(
188                    "rule returned {} cotangents for {} inputs",
189                    cotangent_in.len(),
190                    rule_inputs.len()
191                ),
192            ));
193        }
194
195        for (input, maybe_cotangent) in rule_inputs.iter().zip(cotangent_in) {
196            let cotangent_id = match maybe_cotangent {
197                Some(cotangent_id) => cotangent_id,
198                None => continue,
199            };
200            let input_key = input.key().clone();
201
202            match cotangent_env.get(&input_key).copied() {
203                Some(existing_id) => {
204                    let sum = builder.add_primitive(
205                        Op::add(),
206                        vec![
207                            PrimitiveValue::Local(existing_id),
208                            PrimitiveValue::Local(cotangent_id),
209                        ],
210                        OperationRole::Linearized {
211                            active_mask: vec![true, true],
212                        },
213                    );
214                    cotangent_env.insert(input_key, sum[0]);
215                }
216                None => {
217                    cotangent_env.insert(input_key, cotangent_id);
218                }
219            }
220        }
221    }
222
223    Ok(linear
224        .tangent_inputs()
225        .iter()
226        .map(|(_, tangent_input_id)| {
227            let tangent_input_key = &graph.values()[*tangent_input_id].key;
228            cotangent_env.get(tangent_input_key).copied()
229        })
230        .collect())
231}
232
233fn transpose_input_for<Op: Primitive>(
234    linear: &LinearizedGraph<Op>,
235    input: &ValueRef<Op>,
236) -> PrimitiveTransposeInput<Op>
237where
238    Op::InputKey: ADKey,
239{
240    let graph = linear.as_graph();
241    match input {
242        ValueRef::Local(local_id) => {
243            let key = graph.values()[*local_id].key.clone();
244            PrimitiveTransposeInput::Linear {
245                key,
246                primal: linear.linear_primal(*local_id).cloned(),
247            }
248        }
249        ValueRef::External(key) => PrimitiveTransposeInput::Residual(key.clone()),
250    }
251}
252
253fn cotangent_seed_key<Op: Primitive>(
254    linear: &LinearizedGraph<Op>,
255    index: usize,
256) -> ADRuleResult<Op::InputKey>
257where
258    Op::InputKey: ADKey,
259{
260    if linear.tangent_inputs().is_empty() {
261        return Err(ADRuleError::invalid_input(
262            "tidu::linear_transpose",
263            ADRuleKind::Transpose,
264            "active tangent outputs require at least one tangent input to derive seed keys",
265        ));
266    }
267
268    let base_slot = index.min(linear.tangent_inputs().len() - 1);
269    let base_key = &linear.tangent_inputs()[base_slot].0;
270    Ok(base_key.tangent_of(u64::MAX - index as u64))
271}
tidu/linear_transpose.rs

tidu/
linear_transpose.rs
