tenferro_ad/
extension.rs

1//! Eager AD support for out-of-tree extension primitives.
2
3use std::sync::Arc;
4
5use computegraph::GraphOperation;
6use tenferro_ops::std_tensor_op::StdTensorOp;
7use tenferro_runtime::ad_support::push_metadata_scope;
8use tenferro_runtime::{Error, Result};
9use tenferro_tensor::{Tensor, TensorValue};
10
11use crate::eager::{eager_grad_recording_enabled, record_eager_outputs, EagerRuntime, EagerTensor};
12
13pub use tenferro_ops::ext_op::{
14    ExtensionLinearTransposeRule, ExtensionLinearizeRule, ExtensionPrimalVjpRule,
15    ExtensionRegistryError, ExtensionRuleRole, ExtensionRuleSet, HostReference,
16};
17pub use tenferro_runtime::extension::{
18    apply, ExtensionCacheKey, ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore,
19    ExtensionExecutionContext, ExtensionExecutor, ExtensionFamilyId, ExtensionOp,
20    ExtensionRegistry, ExtensionRuntime, ExtensionRuntimeRegistryError,
21};
22
23/// Adopt an untracked eager tensor value produced by this runtime's backend.
24///
25/// This is a low-level extension contract for eager composite operations that
26/// execute through [`EagerRuntime::with_backend_mut`] and receive a lazy
27/// [`TensorValue`] from the backend. The value must have been produced for the
28/// same eager runtime; this helper intentionally does not register gradient
29/// metadata and must not be used for tracked outputs.
30///
31/// # Examples
32///
33/// ```rust
34/// use tenferro_ad::extension::adopt_untracked_eager_value;
35/// use tenferro_ad::EagerRuntime;
36/// use tenferro_cpu::CpuBackend;
37/// use tenferro_tensor::{Tensor, TensorValue};
38///
39/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
40/// let value = TensorValue::from_tensor(
41///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
42/// );
43/// let eager = adopt_untracked_eager_value(ctx, value);
44/// assert_eq!(eager.shape(), &[1]);
45/// assert!(!eager.tracks_grad());
46/// ```
47#[must_use]
48pub fn adopt_untracked_eager_value(ctx: Arc<EagerRuntime>, value: TensorValue) -> EagerTensor {
49    EagerTensor::new_untracked_value_result(ctx, value)
50}
51
52/// Apply an extension op to eager AD tensors.
53///
54/// # Examples
55///
56/// ```rust
57/// use tenferro_ad::extension::apply_eager;
58/// use tenferro_ad::{EagerRuntime, EagerTensor};
59/// use tenferro_cpu::CpuBackend;
60/// use tenferro_tensor::Tensor;
61///
62/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
63/// let x = EagerTensor::from_tensor_in(
64///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
65///     ctx,
66/// ).unwrap();
67/// let _ = &x;
68/// let _apply = apply_eager;
69/// ```
70pub fn apply_eager(op: Arc<dyn ExtensionOp>, inputs: &[&EagerTensor]) -> Result<Vec<EagerTensor>> {
71    let Some(first) = inputs.first() else {
72        return Err(Error::Internal(
73            "extension::apply_eager requires at least one input tensor".to_string(),
74        ));
75    };
76    if inputs.len() != op.input_count() {
77        return Err(Error::Internal(format!(
78            "extension::apply_eager: op family {:?} expects {} inputs, got {}",
79            op.family_id(),
80            op.input_count(),
81            inputs.len()
82        )));
83    }
84
85    let ctx = Arc::clone(&first.ctx);
86    for tensor in inputs.iter().skip(1) {
87        if !first.same_context(tensor) {
88            return Err(Error::ContextMismatch {
89                lhs: first.ctx_id(),
90                rhs: tensor.ctx_id(),
91            });
92        }
93    }
94
95    let op = StdTensorOp::Extension(op);
96    let input_reads: Vec<_> = inputs.iter().map(|tensor| tensor.tensor_read()).collect();
97    let outputs = ctx.exec_outputs_read(&op, &input_reads)?;
98    if outputs.len() != op.output_count() {
99        return Err(Error::Internal(format!(
100            "expected {} eager outputs for {:?}, got {}",
101            op.output_count(),
102            op,
103            outputs.len()
104        )));
105    }
106
107    if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
108        return outputs
109            .into_iter()
110            .map(|output| EagerTensor::new_untracked_result(Arc::clone(&ctx), output))
111            .collect();
112    }
113
114    let outputs: Vec<Arc<Tensor>> = outputs.into_iter().map(Arc::new).collect();
115    let recorded = record_eager_outputs(&op, &outputs, inputs)?;
116    if recorded.traces.len() != outputs.len() {
117        return Err(Error::Internal(format!(
118            "expected {} eager traces for {:?}, got {}",
119            outputs.len(),
120            op,
121            recorded.traces.len()
122        )));
123    }
124    let mut metadata_scopes = vec![Arc::clone(&recorded.metadata_scope)];
125    for input in inputs {
126        for scope in &input.metadata_scopes {
127            push_metadata_scope(&mut metadata_scopes, Arc::clone(scope));
128        }
129    }
130
131    recorded
132        .traces
133        .into_iter()
134        .zip(outputs)
135        .map(|(trace, output)| {
136            EagerTensor::new_result(
137                Arc::clone(&ctx),
138                trace.key,
139                output.as_ref().clone(),
140                trace.requires_grad,
141                trace.trace,
142                metadata_scopes.clone(),
143            )
144        })
145        .collect()
146}
147
148/// Apply one standard tensor op eagerly and record it for AD when needed.
149///
150/// Extension crates use this when an extension-level eager operation expands
151/// into ordinary `StdTensorOp` nodes instead of a custom extension primitive.
152pub fn apply_standard_op(op: StdTensorOp, inputs: &[&EagerTensor]) -> Result<EagerTensor> {
153    if matches!(op, StdTensorOp::Extension(_)) {
154        return Err(Error::Internal(
155            "extension::apply_standard_op does not accept Extension ops".into(),
156        ));
157    }
158    EagerTensor::nary_op(inputs, op)
159}
tenferro_ad/extension.rs

tenferro_ad/
extension.rs
