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, ErrorPhase, ExtensionModule, Result};
9use tenferro_tensor::{BackendSession, Tensor, TensorRead, TensorValue};
10
11use crate::eager::{eager_grad_recording_enabled, record_eager_outputs, EagerRuntime, EagerTensor};
12
13pub use tenferro_runtime::extension::{
14    apply, ExtensionCacheKey, ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore,
15    ExtensionExecutionContext, ExtensionFamilyId, ExtensionOp,
16};
17
18/// Adopt an untracked eager tensor value produced by this runtime's backend.
19///
20/// This is a low-level extension contract for eager composite operations that
21/// execute through a lifetime-bound backend session and receive a lazy
22/// [`TensorValue`] from the backend. The value must have been produced for the
23/// same eager runtime; this helper intentionally does not register gradient
24/// metadata and must not be used for tracked outputs.
25///
26/// # Examples
27///
28/// ```rust
29/// use tenferro_ad::extension::adopt_untracked_eager_value;
30/// use tenferro_ad::EagerRuntime;
31/// use tenferro_cpu::CpuBackend;
32/// use tenferro_tensor::{Tensor, TensorValue};
33///
34/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
35/// let value = TensorValue::from_tensor(
36///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
37/// );
38/// let eager = adopt_untracked_eager_value(ctx, value);
39/// assert_eq!(eager.shape(), &[1]);
40/// assert!(!eager.tracks_grad());
41/// # Ok::<(), tenferro_ad::Error>(())
42/// ```
43#[must_use]
44pub fn adopt_untracked_eager_value(ctx: Arc<EagerRuntime>, value: TensorValue) -> EagerTensor {
45    EagerTensor::new_untracked_value_result(ctx, value)
46}
47
48/// Apply an extension op to eager AD tensors.
49///
50/// # Examples
51///
52/// ```rust
53/// use tenferro_ad::extension::apply_eager;
54/// use tenferro_ad::{EagerRuntime, EagerTensor};
55/// use tenferro_cpu::CpuBackend;
56/// use tenferro_tensor::Tensor;
57///
58/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
59/// let x = EagerTensor::from_tensor_in(
60///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
61///     ctx,
62/// ).unwrap();
63/// let _ = &x;
64/// let _apply = apply_eager;
65/// # Ok::<(), tenferro_ad::Error>(())
66/// ```
67/// # Errors
68///
69/// Returns `Error::Validation` with `InvalidArgument` when `inputs` is empty
70/// or its length differs from the extension's declared input count. Returns
71/// `Error::ContextMismatch` when tensors belong to different eager runtimes;
72/// backend, extension, and runtime-state failures retain their typed sources.
73pub fn apply_eager(op: Arc<dyn ExtensionOp>, inputs: &[&EagerTensor]) -> Result<Vec<EagerTensor>> {
74    let ctx = validate_eager_extension_inputs(op.as_ref(), inputs)?;
75    let std_op = StdTensorOp::Extension(op);
76    let input_reads: Vec<_> = inputs.iter().map(|tensor| tensor.tensor_read()).collect();
77    let outputs = ctx.exec_outputs_read(&std_op, &input_reads)?;
78    finish_eager_extension_outputs(ctx, std_op, inputs, outputs)
79}
80
81/// Apply an extension op to eager tensors through a direct prepared-operation
82/// callback receiving a non-owning backend session.
83///
84/// This keeps eager extension execution on the extension crate's semantic
85/// implementation without using the retired legacy extension executor.
86///
87/// # Errors
88///
89/// Returns `Error::Validation` with `InvalidArgument` when `inputs` is empty
90/// or its length differs from the extension's declared input count. Returns
91/// `Error::ContextMismatch` when tensors belong to different eager runtimes;
92/// backend, extension, and runtime-state failures retain their typed sources.
93pub fn apply_eager_with_extension_session(
94    op: Arc<dyn ExtensionOp>,
95    inputs: &[&EagerTensor],
96    module: Arc<dyn ExtensionModule>,
97    execute: impl FnOnce(
98            &dyn ExtensionOp,
99            &[TensorRead<'_>],
100            &mut ExtensionExecutionContext<'_, dyn BackendSession + '_>,
101        ) -> tenferro_tensor::Result<Vec<Tensor>>
102        + Send,
103) -> Result<Vec<EagerTensor>> {
104    let ctx = validate_eager_extension_inputs(op.as_ref(), inputs)?;
105    ctx.install_extension_module(module)?;
106    let input_reads: Vec<_> = inputs.iter().map(|tensor| tensor.tensor_read()).collect();
107    let outputs = ctx.with_extension_execution_context(|extension_ctx| {
108        execute(op.as_ref(), &input_reads, extension_ctx)
109    })??;
110    finish_eager_extension_outputs(ctx, StdTensorOp::Extension(op), inputs, outputs)
111}
112
113fn validate_eager_extension_inputs(
114    op: &dyn ExtensionOp,
115    inputs: &[&EagerTensor],
116) -> Result<Arc<EagerRuntime>> {
117    let Some(first) = inputs.first() else {
118        return Err(Error::invalid_argument(
119            "extension::apply_eager",
120            ErrorPhase::Execution,
121            "inputs",
122            "at least one input tensor is required",
123        ));
124    };
125    if inputs.len() != op.input_count() {
126        return Err(Error::invalid_argument(
127            "extension::apply_eager",
128            ErrorPhase::Execution,
129            "inputs",
130            format!(
131                "op family {:?} expects {} inputs, got {}",
132                op.family_id(),
133                op.input_count(),
134                inputs.len()
135            ),
136        ));
137    }
138
139    let ctx = Arc::clone(&first.ctx);
140    for tensor in inputs.iter().skip(1) {
141        if !first.same_context(tensor) {
142            return Err(Error::ContextMismatch {
143                lhs: first.ctx_id(),
144                rhs: tensor.ctx_id(),
145            });
146        }
147    }
148    Ok(ctx)
149}
150
151fn finish_eager_extension_outputs(
152    ctx: Arc<EagerRuntime>,
153    op: StdTensorOp,
154    inputs: &[&EagerTensor],
155    outputs: Vec<Tensor>,
156) -> Result<Vec<EagerTensor>> {
157    if outputs.len() != op.output_count() {
158        return Err(Error::Internal(format!(
159            "expected {} eager outputs for {:?}, got {}",
160            op.output_count(),
161            op,
162            outputs.len()
163        )));
164    }
165
166    if !eager_grad_recording_enabled() {
167        return outputs
168            .into_iter()
169            .map(|output| EagerTensor::new_untracked_result(Arc::clone(&ctx), output))
170            .collect();
171    }
172
173    let outputs: Vec<Arc<Tensor>> = outputs.into_iter().map(Arc::new).collect();
174    let recorded = record_eager_outputs(&op, &outputs, inputs)?;
175    if recorded.traces.len() != outputs.len() {
176        return Err(Error::Internal(format!(
177            "expected {} eager traces for {:?}, got {}",
178            outputs.len(),
179            op,
180            recorded.traces.len()
181        )));
182    }
183    let mut metadata_scopes = vec![Arc::clone(&recorded.metadata_scope)];
184    for input in inputs {
185        for scope in &input.metadata_scopes {
186            push_metadata_scope(&mut metadata_scopes, Arc::clone(scope));
187        }
188    }
189
190    recorded
191        .traces
192        .into_iter()
193        .zip(recorded.semantic_traces)
194        .zip(outputs)
195        .map(|((trace, semantic_trace), output)| {
196            if trace.requires_grad {
197                EagerTensor::new_result_arc_with_semantic_trace(
198                    Arc::clone(&ctx),
199                    trace.key,
200                    output,
201                    trace.requires_grad,
202                    trace.trace,
203                    semantic_trace,
204                    metadata_scopes.clone(),
205                )
206            } else {
207                EagerTensor::new_unregistered_result_arc_with_semantic_trace(
208                    Arc::clone(&ctx),
209                    trace.key,
210                    output,
211                    trace.requires_grad,
212                    trace.trace,
213                    semantic_trace,
214                    metadata_scopes.clone(),
215                )
216            }
217        })
218        .collect()
219}
220
221/// Apply one standard tensor op eagerly and record it for AD when needed.
222///
223/// Extension crates use this when an extension-level eager operation expands
224/// into ordinary `StdTensorOp` nodes instead of a custom extension primitive.
225///
226/// # Errors
227///
228/// Returns [`tenferro_runtime::Error::TensorRuntime`] containing
229/// [`tenferro_tensor::ValidationError::InvalidArgument`] if an extension
230/// op is passed to this standard-op entry point. Returns
231/// [`tenferro_runtime::Error::ContextMismatch`] for tensors from different
232/// eager contexts and propagates typed tensor/backend/runtime-state failures
233/// from the selected eager context.
234pub fn apply_standard_op(op: StdTensorOp, inputs: &[&EagerTensor]) -> Result<EagerTensor> {
235    if matches!(op, StdTensorOp::Extension(_)) {
236        return Err(Error::invalid_argument(
237            "extension::apply_standard_op",
238            ErrorPhase::Execution,
239            "op",
240            "Extension ops must be passed to apply_eager",
241        ));
242    }
243    EagerTensor::nary_op(inputs, op)
244}
tenferro_ad/extension.rs

tenferro_ad/
extension.rs
