tenferro_ad/
traced.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4
5use crate::ad_rule_error::ad_rule_error;
6use computegraph::graph::Graph;
7use computegraph::resolve::resolve;
8use computegraph::resolve::{ResolvedView, ValueDef};
9use computegraph::types::ValueKey;
10use tenferro_ops::input_key::TensorInputKey;
11use tenferro_ops::std_tensor_op::StdTensorOp;
12use tenferro_ops::ExtensionRuleSet;
13use tenferro_ops::ShapeGuardContext;
14use tenferro_runtime::ad_support::{
15    checkpoint_chain as tensor_checkpoint_chain, checkpoint_tensor,
16    extra_roots as tensor_extra_roots, inputs_map as tensor_inputs_map, leaf_input_key,
17    linear_input_key, metadata_scopes as tensor_metadata_scopes, metadata_scopes_with_new,
18    ones_tensor, push_metadata_scope, register_scoped_graph_metadata, registered_meta,
19    resolve_roots as tensor_resolve_roots, shape_hint as tensor_shape_hint, tensor_from_parts,
20    tensor_meta_from_tensor, GlobalMetadataScope, TracedTensorParts,
21};
22use tenferro_runtime::{Error, GraphCompiler, GraphExecutor, Result, TracedTensor};
23use tenferro_tensor::TensorBackend;
24use tidu::{linear_transpose, linearize, ADRuleError};
25
26#[path = "traced/optimizer.rs"]
27mod optimizer;
28#[path = "traced/primal_transpose.rs"]
29mod primal_transpose;
30
31use optimizer::OptimizedLinearGraph;
32use primal_transpose::{try_primal_transpose, PrimalTransposeGraph};
33
34use crate::transform_cache::{
35    AdTransformCache, CachedTracedVjpTransform, TracedAdTransformCacheKey, TracedAdTransformKind,
36};
37
38static NEXT_DIFF_PASS_ID: AtomicU64 = AtomicU64::new(0);
39
40fn next_pass_id() -> u64 {
41    NEXT_DIFF_PASS_ID.fetch_add(1, Ordering::Relaxed)
42}
43
44pub(crate) fn next_input_key() -> TensorInputKey {
45    tenferro_runtime::ad_support::allocate_input_key()
46}
47
48fn error_shape_hint(tensor: &TracedTensor) -> Vec<usize> {
49    tensor
50        .try_concrete_shape()
51        .unwrap_or_else(|| vec![0; tensor.rank])
52}
53
54fn shape_guard_context(
55    extension_rules: Option<&ExtensionRuleSet>,
56    active_values: Option<Arc<HashSet<ValueKey<StdTensorOp>>>>,
57    roots: &[Arc<Graph<StdTensorOp>>],
58) -> ShapeGuardContext {
59    let mut ctx = ShapeGuardContext::with_global_metadata();
60    register_shape_sources(&mut ctx, roots);
61    let ctx = match extension_rules {
62        Some(rules) => ctx.with_extension_rules(rules.clone()),
63        None => ctx,
64    };
65    match active_values {
66        Some(keys) => ctx.with_linearize_active_values(keys),
67        None => ctx,
68    }
69}
70
71fn register_shape_sources(ctx: &mut ShapeGuardContext, roots: &[Arc<Graph<StdTensorOp>>]) {
72    let mut seen = HashSet::new();
73    for graph in roots {
74        register_graph_shape_sources(ctx, graph, &mut seen);
75    }
76}
77
78fn register_graph_shape_sources(
79    ctx: &mut ShapeGuardContext,
80    graph: &Arc<Graph<StdTensorOp>>,
81    seen: &mut HashSet<*const Graph<StdTensorOp>>,
82) {
83    if !seen.insert(Arc::as_ptr(graph)) {
84        return;
85    }
86    for parent in graph.parents() {
87        register_graph_shape_sources(ctx, parent, seen);
88    }
89    for &input_id in graph.inputs() {
90        let key = graph.values()[input_id].key.clone();
91        let Ok(meta) = registered_meta(&key) else {
92            continue;
93        };
94        let Some(shape) = meta.bound_shape() else {
95            continue;
96        };
97        for tensor_id in shape
98            .iter()
99            .flat_map(|dim| dim.referenced_tensor_ids().into_iter())
100        {
101            ctx.insert_shape_source(tensor_id, key.clone());
102        }
103    }
104}
105
106fn linearize_active_value_keys(
107    view: &ResolvedView<StdTensorOp>,
108    outputs: &[ValueKey<StdTensorOp>],
109    aliases: &std::collections::HashMap<TensorInputKey, ValueKey<StdTensorOp>>,
110) -> Arc<HashSet<ValueKey<StdTensorOp>>> {
111    let mut active = HashSet::new();
112    let mut stack: Vec<ValueKey<StdTensorOp>> = outputs.to_vec();
113    while let Some(key) = stack.pop() {
114        if !active.insert(key.clone()) {
115            continue;
116        }
117        let Some(val_def) = view.resolve_value(&key) else {
118            continue;
119        };
120        match val_def {
121            ValueDef::Produced { input_keys, .. } => {
122                for input_key in input_keys {
123                    stack.push(input_key.clone());
124                }
125            }
126            ValueDef::Input { key: input_key } => {
127                if let Some(aliased) = aliases.get(&input_key) {
128                    stack.push(aliased.clone());
129                }
130            }
131        }
132    }
133    Arc::new(active)
134}
135
136fn graph_has_registered_primal_vjp(
137    view: &ResolvedView<StdTensorOp>,
138    outputs: &[ValueKey<StdTensorOp>],
139    aliases: &HashMap<TensorInputKey, ValueKey<StdTensorOp>>,
140    extension_rules: Option<&ExtensionRuleSet>,
141) -> bool {
142    let Some(extension_rules) = extension_rules else {
143        return false;
144    };
145    let mut seen = HashSet::new();
146    let mut stack = outputs.to_vec();
147    while let Some(key) = stack.pop() {
148        if !seen.insert(key.clone()) {
149            continue;
150        }
151        if let ValueKey::Derived { operation, .. } = &key {
152            if let StdTensorOp::Extension(ext) = operation.operation() {
153                if extension_rules.lookup_primal_vjp(ext.family_id()).is_some() {
154                    return true;
155                }
156            }
157        }
158        let Some(val_def) = view.resolve_value(&key) else {
159            continue;
160        };
161        match val_def {
162            ValueDef::Produced { input_keys, .. } => {
163                for input_key in input_keys {
164                    stack.push(input_key);
165                }
166            }
167            ValueDef::Input { key: input_key } => {
168                if let Some(aliased) = aliases.get(&input_key) {
169                    stack.push(aliased.clone());
170                }
171            }
172        }
173    }
174    false
175}
176
177fn is_not_applicable_custom_vjp(err: &ADRuleError) -> bool {
178    matches!(err, ADRuleError::Unsupported { .. })
179}
180
181pub(crate) fn grad_with_rules_and_cache(
182    output: &TracedTensor,
183    wrt: &TracedTensor,
184    extension_rules: &ExtensionRuleSet,
185    ad_transform_cache: Option<&AdTransformCache>,
186) -> Result<TracedTensor> {
187    grad_with_optional_rules(output, wrt, Some(extension_rules), ad_transform_cache)
188}
189
190pub(crate) fn jvp_with_rules_and_cache(
191    output: &TracedTensor,
192    wrt: &TracedTensor,
193    tangent: &TracedTensor,
194    extension_rules: &ExtensionRuleSet,
195    ad_transform_cache: Option<&AdTransformCache>,
196) -> Result<TracedTensor> {
197    let wrt_input_key = leaf_input_key(wrt)?;
198    jvp_optional_impl(
199        output,
200        wrt,
201        tangent,
202        Some(extension_rules),
203        ad_transform_cache,
204    )?
205    .ok_or_else(|| Error::Internal(format!("jvp output is inactive for {:?}", wrt_input_key)))
206}
207
208pub(crate) fn grad_optional_with_rules_and_cache(
209    output: &TracedTensor,
210    wrt: &TracedTensor,
211    extension_rules: &ExtensionRuleSet,
212    ad_transform_cache: Option<&AdTransformCache>,
213) -> Result<Option<TracedTensor>> {
214    if output.rank != 0 {
215        return Err(Error::NonScalarGrad {
216            shape: error_shape_hint(output),
217        });
218    }
219
220    let ones = ones_tensor(output.dtype, vec![])?;
221    let seed = TracedTensor::from_tensor_concrete_shape(ones)?;
222    vjp_optional_impl(
223        output,
224        wrt,
225        &seed,
226        Some(extension_rules),
227        "grad",
228        ad_transform_cache,
229    )
230}
231
232pub(crate) fn jvp_optional_with_rules_and_cache(
233    output: &TracedTensor,
234    wrt: &TracedTensor,
235    tangent: &TracedTensor,
236    extension_rules: &ExtensionRuleSet,
237    ad_transform_cache: Option<&AdTransformCache>,
238) -> Result<Option<TracedTensor>> {
239    jvp_optional_impl(
240        output,
241        wrt,
242        tangent,
243        Some(extension_rules),
244        ad_transform_cache,
245    )
246}
247
248pub(crate) fn vjp_with_rules_and_cache(
249    output: &TracedTensor,
250    wrt: &TracedTensor,
251    cotangent: &TracedTensor,
252    extension_rules: &ExtensionRuleSet,
253    ad_transform_cache: Option<&AdTransformCache>,
254) -> Result<TracedTensor> {
255    let wrt_input_key = leaf_input_key(wrt)?;
256    vjp_optional_impl(
257        output,
258        wrt,
259        cotangent,
260        Some(extension_rules),
261        "vjp",
262        ad_transform_cache,
263    )?
264    .ok_or_else(|| Error::Internal(format!("vjp output is inactive for {:?}", wrt_input_key)))
265}
266
267pub(crate) fn vjp_optional_with_rules_and_cache(
268    output: &TracedTensor,
269    wrt: &TracedTensor,
270    cotangent: &TracedTensor,
271    extension_rules: &ExtensionRuleSet,
272    ad_transform_cache: Option<&AdTransformCache>,
273) -> Result<Option<TracedTensor>> {
274    vjp_optional_impl(
275        output,
276        wrt,
277        cotangent,
278        Some(extension_rules),
279        "vjp",
280        ad_transform_cache,
281    )
282}
283
284fn grad_with_optional_rules(
285    output: &TracedTensor,
286    wrt: &TracedTensor,
287    extension_rules: Option<&ExtensionRuleSet>,
288    ad_transform_cache: Option<&AdTransformCache>,
289) -> Result<TracedTensor> {
290    if output.rank != 0 {
291        return Err(Error::NonScalarGrad {
292            shape: error_shape_hint(output),
293        });
294    }
295
296    let ones = ones_tensor(output.dtype, vec![])?;
297    let seed = TracedTensor::from_tensor_concrete_shape(ones)?;
298    let wrt_input_key = leaf_input_key(wrt)?;
299    vjp_optional_impl(
300        output,
301        wrt,
302        &seed,
303        extension_rules,
304        "grad",
305        ad_transform_cache,
306    )?
307    .ok_or_else(|| Error::Internal(format!("grad output is inactive for {:?}", wrt_input_key)))
308}
309
310/// Automatic differentiation helpers for [`TracedTensor`].
311///
312/// # Examples
313///
314/// ```rust
315/// use tenferro_ad::TracedTensorAdExt;
316/// use tenferro_runtime::TracedTensor;
317///
318/// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
319/// let loss = x.scale_real(2.0).unwrap();
320/// let maybe_dx = loss.grad_optional(&x).unwrap();
321/// assert!(maybe_dx.is_some());
322/// ```
323pub trait TracedTensorAdExt {
324    /// Gradient of a scalar output with respect to a traced input.
325    ///
326    /// For complex scalar outputs, tenferro returns the Hermitian-adjoint
327    /// cotangent. To compare seed-`1` scalar gradients with JAX's public
328    /// `grad` values, use the complex conjugate of this result. See
329    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
330    ///
331    /// # Examples
332    ///
333    /// ```rust
334    /// use tenferro_ad::TracedTensorAdExt;
335    /// use tenferro_cpu::CpuBackend;
336    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, TracedTensor};
337    ///
338    /// fn eval(tensor: &TracedTensor) -> tenferro_runtime::Tensor {
339    ///     let mut compiler = GraphCompiler::new();
340    ///     let program = compiler.compile(tensor).unwrap();
341    ///     let mut executor = GraphExecutor::new(CpuBackend::new());
342    ///     executor.run(&program).unwrap()
343    /// }
344    ///
345    /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
346    /// let loss = (&x * &x).unwrap();
347    /// let dx = loss.grad(&x).unwrap();
348    ///
349    /// assert_eq!(eval(&dx).as_slice::<f64>().unwrap(), &[6.0]);
350    /// ```
351    fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor>;
352
353    /// Like [`grad`](Self::grad), but returns `None` when `wrt` is inactive.
354    ///
355    /// # Examples
356    ///
357    /// ```rust
358    /// use tenferro_ad::TracedTensorAdExt;
359    /// use tenferro_runtime::TracedTensor;
360    ///
361    /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
362    /// let y = TracedTensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap();
363    /// let loss = (&y * &y).unwrap();
364    ///
365    /// assert!(loss.grad_optional(&x).unwrap().is_none());
366    /// ```
367    fn grad_optional(&self, wrt: &TracedTensor) -> Result<Option<TracedTensor>>;
368
369    /// Evaluate this tensor and replace its graph with a concrete leaf while
370    /// preserving the previous graph for AD replay.
371    ///
372    /// # Examples
373    ///
374    /// ```rust
375    /// use tenferro_ad::TracedTensorAdExt;
376    /// use tenferro_cpu::CpuBackend;
377    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, TracedTensor};
378    ///
379    /// let mut compiler = GraphCompiler::new();
380    /// let mut executor = GraphExecutor::new(CpuBackend::new());
381    /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
382    /// let mut y = (&x * &x).unwrap();
383    ///
384    /// y.checkpoint(&mut compiler, &mut executor).unwrap();
385    ///
386    /// let value = y.attached_data().unwrap();
387    /// assert_eq!(value.as_slice::<f64>().unwrap(), &[9.0]);
388    /// ```
389    fn checkpoint<B: TensorBackend>(
390        &mut self,
391        compiler: &mut GraphCompiler,
392        executor: &mut GraphExecutor<B>,
393    ) -> Result<()>;
394
395    /// Forward-mode Jacobian-vector product.
396    ///
397    /// # Examples
398    ///
399    /// ```rust
400    /// use tenferro_ad::TracedTensorAdExt;
401    /// use tenferro_cpu::CpuBackend;
402    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, TracedTensor};
403    ///
404    /// fn eval(tensor: &TracedTensor) -> tenferro_runtime::Tensor {
405    ///     let mut compiler = GraphCompiler::new();
406    ///     let program = compiler.compile(tensor).unwrap();
407    ///     let mut executor = GraphExecutor::new(CpuBackend::new());
408    ///     executor.run(&program).unwrap()
409    /// }
410    ///
411    /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
412    /// let tangent = TracedTensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap();
413    /// let y = (&x * &x).unwrap();
414    /// let dy = y.jvp(&x, &tangent).unwrap();
415    ///
416    /// assert_eq!(eval(&dy).as_slice::<f64>().unwrap(), &[12.0]);
417    /// ```
418    fn jvp(&self, wrt: &TracedTensor, tangent: &TracedTensor) -> Result<TracedTensor>;
419
420    /// Like [`jvp`](Self::jvp), but returns `None` when `wrt` is inactive.
421    ///
422    /// # Examples
423    ///
424    /// ```rust
425    /// use tenferro_ad::TracedTensorAdExt;
426    /// use tenferro_runtime::TracedTensor;
427    ///
428    /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
429    /// let y = TracedTensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap();
430    /// let tangent = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
431    /// let loss = (&y * &y).unwrap();
432    ///
433    /// assert!(loss.jvp_optional(&x, &tangent).unwrap().is_none());
434    /// ```
435    fn jvp_optional(
436        &self,
437        wrt: &TracedTensor,
438        tangent: &TracedTensor,
439    ) -> Result<Option<TracedTensor>>;
440
441    /// Reverse-mode vector-Jacobian product.
442    ///
443    /// Complex cotangents use tenferro's Hermitian real-inner-product
444    /// convention. Non-real complex cotangent seeds therefore need an explicit
445    /// seed-convention comparison when matching JAX. See
446    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
447    ///
448    /// # Examples
449    ///
450    /// ```rust
451    /// use tenferro_ad::TracedTensorAdExt;
452    /// use tenferro_cpu::CpuBackend;
453    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, TracedTensor};
454    ///
455    /// fn eval(tensor: &TracedTensor) -> tenferro_runtime::Tensor {
456    ///     let mut compiler = GraphCompiler::new();
457    ///     let program = compiler.compile(tensor).unwrap();
458    ///     let mut executor = GraphExecutor::new(CpuBackend::new());
459    ///     executor.run(&program).unwrap()
460    /// }
461    ///
462    /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
463    /// let cotangent = TracedTensor::from_vec_col_major(vec![], vec![0.5_f64]).unwrap();
464    /// let y = (&x * &x).unwrap();
465    /// let dx = y.vjp(&x, &cotangent).unwrap();
466    ///
467    /// assert_eq!(eval(&dx).as_slice::<f64>().unwrap(), &[3.0]);
468    /// ```
469    fn vjp(&self, wrt: &TracedTensor, cotangent: &TracedTensor) -> Result<TracedTensor>;
470
471    /// Like [`vjp`](Self::vjp), but returns `None` when `wrt` is inactive.
472    ///
473    /// # Examples
474    ///
475    /// ```rust
476    /// use tenferro_ad::TracedTensorAdExt;
477    /// use tenferro_runtime::TracedTensor;
478    ///
479    /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
480    /// let y = TracedTensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap();
481    /// let cotangent = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
482    /// let loss = (&y * &y).unwrap();
483    ///
484    /// assert!(loss.vjp_optional(&x, &cotangent).unwrap().is_none());
485    /// ```
486    fn vjp_optional(
487        &self,
488        wrt: &TracedTensor,
489        cotangent: &TracedTensor,
490    ) -> Result<Option<TracedTensor>>;
491}
492
493impl TracedTensorAdExt for TracedTensor {
494    fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor> {
495        grad_with_optional_rules(self, wrt, None, None)
496    }
497
498    fn grad_optional(&self, wrt: &TracedTensor) -> Result<Option<TracedTensor>> {
499        if self.rank != 0 {
500            return Err(Error::NonScalarGrad {
501                shape: error_shape_hint(self),
502            });
503        }
504
505        let ones = ones_tensor(self.dtype, vec![])?;
506        let seed = TracedTensor::from_tensor_concrete_shape(ones)?;
507        vjp_optional_impl(self, wrt, &seed, None, "grad", None)
508    }
509
510    fn checkpoint<B: TensorBackend>(
511        &mut self,
512        compiler: &mut GraphCompiler,
513        executor: &mut GraphExecutor<B>,
514    ) -> Result<()> {
515        let data = if let Some(data) = self.attached_data() {
516            Arc::clone(data)
517        } else {
518            let program = compiler.compile(self)?;
519            Arc::new(executor.run(&program)?)
520        };
521        checkpoint_tensor(self, data)?;
522        Ok(())
523    }
524
525    fn jvp(&self, wrt: &TracedTensor, tangent: &TracedTensor) -> Result<TracedTensor> {
526        let wrt_input_key = leaf_input_key(wrt)?;
527        self.jvp_optional(wrt, tangent)?.ok_or_else(|| {
528            Error::Internal(format!("jvp output is inactive for {:?}", wrt_input_key))
529        })
530    }
531
532    fn jvp_optional(
533        &self,
534        wrt: &TracedTensor,
535        tangent: &TracedTensor,
536    ) -> Result<Option<TracedTensor>> {
537        jvp_optional_impl(self, wrt, tangent, None, None)
538    }
539
540    fn vjp(&self, wrt: &TracedTensor, cotangent: &TracedTensor) -> Result<TracedTensor> {
541        let wrt_input_key = leaf_input_key(wrt)?;
542        self.vjp_optional(wrt, cotangent)?.ok_or_else(|| {
543            Error::Internal(format!("vjp output is inactive for {:?}", wrt_input_key))
544        })
545    }
546
547    fn vjp_optional(
548        &self,
549        wrt: &TracedTensor,
550        cotangent: &TracedTensor,
551    ) -> Result<Option<TracedTensor>> {
552        vjp_optional_impl(self, wrt, cotangent, None, "vjp", None)
553    }
554}
555
556fn jvp_optional_impl(
557    output: &TracedTensor,
558    wrt: &TracedTensor,
559    tangent: &TracedTensor,
560    extension_rules: Option<&ExtensionRuleSet>,
561    ad_transform_cache: Option<&AdTransformCache>,
562) -> Result<Option<TracedTensor>> {
563    let wrt_input_key = leaf_input_key(wrt)?;
564    let output_key = output.graph().values()[output.val].key.clone();
565    let checkpoint_chain = tensor_checkpoint_chain(output);
566    let aliases = checkpoint_chain
567        .as_ref()
568        .map(|chain| chain.collect_aliases())
569        .unwrap_or_default();
570    let checkpoint_graphs = checkpoint_chain
571        .as_ref()
572        .map(|chain| chain.collect_graphs())
573        .unwrap_or_default();
574    let mut roots = tensor_resolve_roots(output);
575    roots.extend(checkpoint_graphs.iter().cloned());
576    let view = resolve(roots);
577    let active_values =
578        linearize_active_value_keys(&view, std::slice::from_ref(&output_key), &aliases);
579    let cache_key = ad_transform_cache.map(|_| {
580        TracedAdTransformCacheKey::new(
581            TracedAdTransformKind::Jvp,
582            &view.roots,
583            &output_key,
584            &wrt_input_key,
585            &aliases,
586        )
587    });
588    let linear = match (ad_transform_cache, cache_key) {
589        (Some(cache), Some(key)) => {
590            if let Some(linear) = cache.get_traced_linearized(&key)? {
591                linear
592            } else {
593                let mut ad_ctx =
594                    shape_guard_context(extension_rules, Some(active_values), &view.roots);
595                let linear = linearize(
596                    &view,
597                    std::slice::from_ref(&output_key),
598                    std::slice::from_ref(&wrt_input_key),
599                    next_pass_id(),
600                    &mut ad_ctx,
601                    &aliases,
602                )
603                .map_err(|err| ad_rule_error("jvp", err))?;
604                let linear = Arc::new(OptimizedLinearGraph::from_tidu(linear).into_cached());
605                cache.put_traced_linearized(key, Arc::clone(&linear))?;
606                linear
607            }
608        }
609        _ => {
610            let mut ad_ctx = shape_guard_context(extension_rules, Some(active_values), &view.roots);
611            let linear = linearize(
612                &view,
613                std::slice::from_ref(&output_key),
614                std::slice::from_ref(&wrt_input_key),
615                next_pass_id(),
616                &mut ad_ctx,
617                &aliases,
618            )
619            .map_err(|err| ad_rule_error("jvp", err))?;
620            Arc::new(OptimizedLinearGraph::from_tidu(linear).into_cached())
621        }
622    };
623    let Some(tangent_output) = linear.tangent_outputs()[0] else {
624        return Ok(None);
625    };
626    let tangent_input_key = linear_input_key(linear.as_graph(), linear.tangent_inputs()[0].1)?;
627    let tangent_data =
628        tangent
629            .attached_data()
630            .cloned()
631            .ok_or_else(|| Error::InvalidGraphBuild {
632                op: "jvp",
633                message: "jvp tangent must have concrete tensor data".to_string(),
634            })?;
635    let metadata_scope = register_scoped_graph_metadata(
636        linear.as_graph(),
637        vec![(
638            ValueKey::Input(tangent_input_key.clone()),
639            tensor_meta_from_tensor(tangent_data.as_ref()),
640        )],
641    )?;
642
643    let mut inputs_map = (*tensor_inputs_map(output)).clone();
644    if let Some(chain) = &checkpoint_chain {
645        inputs_map.extend(chain.collect_inputs());
646    }
647    inputs_map.insert(tangent_input_key, tangent_data);
648
649    let mut extra_roots = vec![Arc::clone(output.graph())];
650    extra_roots.extend(checkpoint_graphs);
651    extra_roots.extend(tensor_extra_roots(output));
652
653    Ok(Some(tensor_from_parts(TracedTensorParts {
654        rank: output.rank,
655        dtype: output.dtype,
656        graph: Arc::clone(linear.graph()),
657        val: tangent_output,
658        data: None,
659        shape_hint: tensor_shape_hint(output),
660        inputs_map: Arc::new(inputs_map),
661        extra_roots,
662        checkpoint_chain,
663        metadata_scopes: metadata_scopes_with_new(
664            metadata_scope,
665            [
666                tensor_metadata_scopes(output),
667                tensor_metadata_scopes(wrt),
668                tensor_metadata_scopes(tangent),
669            ],
670        ),
671    })))
672}
673
674enum VjpTransposeGraph {
675    Primal(PrimalTransposeGraph),
676    Linear(Arc<CachedTracedVjpTransform>),
677}
678
679struct ActiveLinearVjp {
680    transposed: Arc<CachedTracedVjpTransform>,
681    metadata_scope: GlobalMetadataScope,
682}
683
684impl VjpTransposeGraph {
685    fn as_graph(&self) -> &computegraph::graph::Graph<StdTensorOp> {
686        match self {
687            Self::Primal(graph) => graph.as_graph(),
688            Self::Linear(graph) => graph.transposed().as_graph(),
689        }
690    }
691
692    fn tangent_inputs(&self) -> &[(TensorInputKey, computegraph::LocalValueId)] {
693        match self {
694            Self::Primal(graph) => graph.tangent_inputs(),
695            Self::Linear(graph) => graph.transposed().tangent_inputs(),
696        }
697    }
698
699    fn tangent_outputs(&self) -> &[Option<computegraph::LocalValueId>] {
700        match self {
701            Self::Primal(graph) => graph.tangent_outputs(),
702            Self::Linear(graph) => graph.transposed().tangent_outputs(),
703        }
704    }
705
706    fn into_graph_arc(self) -> Arc<computegraph::graph::Graph<StdTensorOp>> {
707        match self {
708            Self::Primal(graph) => Arc::new(graph.into_graph()),
709            Self::Linear(graph) => Arc::clone(graph.transposed().graph()),
710        }
711    }
712}
713
714fn compute_linear_vjp_transform(
715    view: &ResolvedView<StdTensorOp>,
716    output_key: &ValueKey<StdTensorOp>,
717    wrt_input_key: &TensorInputKey,
718    aliases: &HashMap<TensorInputKey, ValueKey<StdTensorOp>>,
719    extension_rules: Option<&ExtensionRuleSet>,
720    active_values: Arc<HashSet<ValueKey<StdTensorOp>>>,
721    wrt: &TracedTensor,
722) -> Result<tidu::ADRuleResult<Option<ActiveLinearVjp>>> {
723    let mut linear_ad_ctx = shape_guard_context(extension_rules, Some(active_values), &view.roots);
724    let linear = match linearize(
725        view,
726        std::slice::from_ref(output_key),
727        std::slice::from_ref(wrt_input_key),
728        next_pass_id(),
729        &mut linear_ad_ctx,
730        aliases,
731    ) {
732        Ok(linear) => linear,
733        Err(err) => return Ok(Err(err)),
734    };
735    if linear.tangent_outputs()[0].is_none() {
736        return Ok(Ok(None));
737    }
738
739    let linear_seed_key = linear_input_key(linear.as_graph(), linear.tangent_inputs()[0].1)?;
740    let linear_metadata_scope = register_scoped_graph_metadata(
741        linear.as_graph(),
742        vec![(
743            ValueKey::Input(linear_seed_key),
744            registered_meta(&wrt.graph().values()[wrt.val].key)?,
745        )],
746    )?;
747    linear_ad_ctx.refresh_global_metadata();
748    let transposed = match linear_transpose(&linear, &mut linear_ad_ctx) {
749        Ok(transposed) => OptimizedLinearGraph::from_tidu(transposed).into_cached(),
750        Err(err) => return Ok(Err(err)),
751    };
752    let (_linear_graph, residual_graph) = linear.into_graphs();
753    let residual_metadata_scope =
754        register_scoped_graph_metadata(residual_graph.as_ref(), std::iter::empty())?;
755    drop(linear_metadata_scope);
756    Ok(Ok(Some(ActiveLinearVjp {
757        transposed: Arc::new(CachedTracedVjpTransform::new(residual_graph, transposed)),
758        metadata_scope: residual_metadata_scope,
759    })))
760}
761
762fn vjp_optional_impl(
763    output: &TracedTensor,
764    wrt: &TracedTensor,
765    cotangent: &TracedTensor,
766    extension_rules: Option<&ExtensionRuleSet>,
767    transform: &'static str,
768    ad_transform_cache: Option<&AdTransformCache>,
769) -> Result<Option<TracedTensor>> {
770    let wrt_input_key = leaf_input_key(wrt)?;
771    let output_key = output.graph().values()[output.val].key.clone();
772    let checkpoint_chain = tensor_checkpoint_chain(output);
773    let aliases = checkpoint_chain
774        .as_ref()
775        .map(|chain| chain.collect_aliases())
776        .unwrap_or_default();
777    let checkpoint_graphs = checkpoint_chain
778        .as_ref()
779        .map(|chain| chain.collect_graphs())
780        .unwrap_or_default();
781    let mut roots = tensor_resolve_roots(output);
782    roots.extend(checkpoint_graphs.iter().cloned());
783    let view = resolve(roots);
784
785    let active_values =
786        linearize_active_value_keys(&view, std::slice::from_ref(&output_key), &aliases);
787    let cache_key = ad_transform_cache.map(|_| {
788        TracedAdTransformCacheKey::new(
789            TracedAdTransformKind::Vjp,
790            &view.roots,
791            &output_key,
792            &wrt_input_key,
793            &aliases,
794        )
795    });
796    if graph_has_registered_primal_vjp(
797        &view,
798        std::slice::from_ref(&output_key),
799        &aliases,
800        extension_rules,
801    ) {
802        let mut primal_ad_ctx = shape_guard_context(extension_rules, None, &view.roots);
803        primal_ad_ctx.refresh_global_metadata();
804        match try_primal_transpose(
805            &view,
806            std::slice::from_ref(&output_key),
807            std::slice::from_ref(&wrt_input_key),
808            &aliases,
809            &mut primal_ad_ctx,
810            next_pass_id(),
811        ) {
812            Ok(transposed) => {
813                if transposed
814                    .tangent_outputs()
815                    .first()
816                    .and_then(|slot| *slot)
817                    .is_some()
818                {
819                    let transposed = VjpTransposeGraph::Primal(transposed);
820                    return build_vjp_tensor(
821                        output,
822                        wrt,
823                        cotangent,
824                        transposed,
825                        None,
826                        checkpoint_chain,
827                        checkpoint_graphs,
828                    );
829                }
830                return Ok(None);
831            }
832            Err(err) if !is_not_applicable_custom_vjp(&err) => {
833                return Err(ad_rule_error(transform, err));
834            }
835            Err(_) => {}
836        }
837    }
838
839    let linear_attempt = match (ad_transform_cache, cache_key) {
840        (Some(cache), Some(key)) => {
841            if let Some(cached) = cache.get_traced_vjp(&key)? {
842                let residual_metadata_scope =
843                    register_scoped_graph_metadata(cached.residual_graph(), std::iter::empty())?;
844                Ok(Some(ActiveLinearVjp {
845                    transposed: cached,
846                    metadata_scope: residual_metadata_scope,
847                }))
848            } else {
849                let computed = compute_linear_vjp_transform(
850                    &view,
851                    &output_key,
852                    &wrt_input_key,
853                    &aliases,
854                    extension_rules,
855                    active_values,
856                    wrt,
857                )?;
858                if let Ok(Some(active)) = &computed {
859                    cache.put_traced_vjp(key, Arc::clone(&active.transposed))?;
860                }
861                computed
862            }
863        }
864        _ => compute_linear_vjp_transform(
865            &view,
866            &output_key,
867            &wrt_input_key,
868            &aliases,
869            extension_rules,
870            active_values,
871            wrt,
872        )?,
873    };
874
875    let (transposed, linear_metadata_scope) = match linear_attempt {
876        Ok(None) => return Ok(None),
877        Ok(Some(active)) => (
878            VjpTransposeGraph::Linear(active.transposed),
879            Some(active.metadata_scope),
880        ),
881        Err(linear_err) => return Err(ad_rule_error(transform, linear_err)),
882    };
883
884    build_vjp_tensor(
885        output,
886        wrt,
887        cotangent,
888        transposed,
889        linear_metadata_scope,
890        checkpoint_chain,
891        checkpoint_graphs,
892    )
893}
894
895fn build_vjp_tensor(
896    output: &TracedTensor,
897    wrt: &TracedTensor,
898    cotangent: &TracedTensor,
899    transposed: VjpTransposeGraph,
900    linear_metadata_scope: Option<GlobalMetadataScope>,
901    checkpoint_chain: Option<Arc<tenferro_runtime::ad_support::CheckpointNode>>,
902    checkpoint_graphs: Vec<Arc<Graph<StdTensorOp>>>,
903) -> Result<Option<TracedTensor>> {
904    let cotangent_input_key =
905        linear_input_key(transposed.as_graph(), transposed.tangent_inputs()[0].1)?;
906    let cotangent_data =
907        cotangent
908            .attached_data()
909            .cloned()
910            .ok_or_else(|| Error::InvalidGraphBuild {
911                op: "vjp",
912                message: "vjp cotangent must have concrete tensor data".to_string(),
913            })?;
914    let transposed_metadata_scope = register_scoped_graph_metadata(
915        transposed.as_graph(),
916        vec![(
917            ValueKey::Input(cotangent_input_key.clone()),
918            tensor_meta_from_tensor(cotangent_data.as_ref()),
919        )],
920    )?;
921    let Some(cotangent_output) = transposed.tangent_outputs()[0] else {
922        return Ok(None);
923    };
924
925    let mut inputs_map = (*tensor_inputs_map(output)).clone();
926    if let Some(chain) = &checkpoint_chain {
927        inputs_map.extend(chain.collect_inputs());
928    }
929    inputs_map.insert(cotangent_input_key.clone(), cotangent_data);
930
931    let mut extra_roots = vec![Arc::clone(output.graph())];
932    if let VjpTransposeGraph::Linear(cached) = &transposed {
933        extra_roots.push(Arc::clone(cached.residual_graph()));
934    }
935    extra_roots.extend(checkpoint_graphs);
936    extra_roots.extend(tensor_extra_roots(output));
937
938    Ok(Some(tensor_from_parts(TracedTensorParts {
939        rank: wrt.rank,
940        dtype: wrt.dtype,
941        graph: transposed.into_graph_arc(),
942        val: cotangent_output,
943        data: None,
944        shape_hint: tensor_shape_hint(wrt),
945        inputs_map: Arc::new(inputs_map),
946        extra_roots,
947        checkpoint_chain,
948        metadata_scopes: {
949            let mut scopes = if let Some(scope) = linear_metadata_scope {
950                metadata_scopes_with_new(
951                    scope,
952                    [
953                        tensor_metadata_scopes(output),
954                        tensor_metadata_scopes(wrt),
955                        tensor_metadata_scopes(cotangent),
956                    ],
957                )
958            } else {
959                let mut scopes: Vec<Arc<crate::metadata::GlobalMetadataScope>> = Vec::new();
960                for inherited in [
961                    tensor_metadata_scopes(output),
962                    tensor_metadata_scopes(wrt),
963                    tensor_metadata_scopes(cotangent),
964                ] {
965                    for scope in inherited {
966                        scopes.push(Arc::clone(scope));
967                    }
968                }
969                scopes
970            };
971            push_metadata_scope(&mut scopes, Arc::new(transposed_metadata_scope));
972            scopes
973        },
974    })))
975}
tenferro_ad/traced.rs

tenferro_ad/
traced.rs
