Skip to main content

tenferro_ad/
eager.rs

1use std::borrow::Cow;
2use std::cell::{Cell, RefCell};
3use std::cmp::Reverse;
4use std::collections::HashMap;
5use std::env;
6use std::fmt;
7use std::marker::PhantomData;
8use std::mem::{size_of, size_of_val};
9use std::rc::Rc;
10#[cfg(test)]
11use std::sync::atomic::{AtomicUsize, Ordering};
12use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
13use std::time::{Duration, Instant};
14
15use lru::LruCache;
16
17use crate::extension::{
18    validate_eager_extension_target, EagerExtensionBackendKind, EagerExtensionTarget,
19};
20use crate::extension_cache::{ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore};
21#[cfg(test)]
22use computegraph::graph::Graph;
23use computegraph::ValueKey;
24#[cfg(test)]
25use computegraph::ValueRef;
26use tenferro_cpu::{CpuBackend, CpuBackendError, CpuPlacement};
27#[cfg(feature = "cuda")]
28use tenferro_gpu::cuda::CudaBackend;
29#[cfg(feature = "webgpu")]
30use tenferro_gpu::webgpu::WebGpuBackend;
31#[cfg(test)]
32use tenferro_ops::input_key::TensorInputKey;
33use tenferro_ops::{std_tensor_op::StdTensorOp, SymDim, TensorMeta};
34use tenferro_runtime::ad_support::{
35    analyze_deferred_semantic_trace, compile_ad_source, ones_tensor, RetainedValue,
36};
37use tenferro_runtime::program::{ProgramValueMetadata, SemanticFingerprint, SemanticProgram};
38use tenferro_runtime::{
39    CompiledGraph, CoreCapabilityBundle, EngineId, ErrorPhase, ExecutionContextIdentity,
40    ExtensionModule, GraphCompiler, HardwareClassId, PreparedCompiledGraph, RegistrationIdentity,
41    Runtime, RuntimeConfigError, RuntimeConfigSnapshot, RuntimeEpoch, TracedTensor,
42};
43#[cfg(test)]
44use tenferro_tensor::TypedTensor;
45use tenferro_tensor::{
46    AllocationGroup, CacheStats, DType, DescriptorSlot, GroupError, IntoShapeVec, Tensor,
47    TensorBackend, TensorRead, TensorScalar, TensorValue, TensorView,
48};
49use tenferro_tensor::{BackendSession, BackendSessionHost};
50
51#[cfg(feature = "cuda")]
52use crate::eager_backend::cuda_runtime_engine_id;
53use crate::eager_backend::{
54    cpu_runtime_engine_id, cpu_runtime_hardware_class, eager_runtime_for_backend, EagerBackend,
55};
56#[cfg(test)]
57use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
58use crate::eager_exec::{
59    eager_input_promotion_plan, exec_op_on_tensor_reads_with_runtime,
60    exec_op_on_tensors_with_runtime,
61};
62use crate::error::{ContextId, Error, Result};
63use crate::metadata::tensor_meta_from_tensor;
64use crate::semantic_extension::SemanticExtensionRuleSet;
65use crate::traced::{derivative_trace_from_frozen_program, next_input_key};
66use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
67
68use crate::AdContext;
69
70pub(crate) type GradSlot = Arc<Mutex<Option<Arc<AdValueRecord>>>>;
71pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<AdValueRecord>>>>;
72
73#[derive(Clone, Debug)]
74pub(crate) struct EagerTrace;
75
76#[cfg(test)]
77pub(crate) static CPU_RUNTIME_SELECTION_REFRESHES: AtomicUsize = AtomicUsize::new(0);
78
79struct CpuRuntimeSelection {
80    snapshot: Arc<RuntimeConfigSnapshot>,
81    epoch: RuntimeEpoch,
82    engine_id: EngineId,
83    registration_identity: RegistrationIdentity,
84    capabilities: CoreCapabilityBundle,
85}
86
87#[derive(Debug, Default, Clone)]
88struct EagerOpProfileEntry {
89    calls: usize,
90    total_time: Duration,
91}
92
93thread_local! {
94    static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
95        RefCell::new(HashMap::new());
96    static EAGER_NO_GRAD_DEPTH: Cell<usize> = const { Cell::new(0) };
97    static EAGER_CAPTURE_DEPTH: Cell<usize> = const { Cell::new(0) };
98    #[cfg(test)]
99    static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
100    #[cfg(test)]
101    static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
102    #[cfg(test)]
103    static EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
104}
105
106#[cfg(test)]
107pub(crate) static EAGER_SEMANTIC_VJP_EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
108
109pub(crate) fn eager_grad_recording_enabled() -> bool {
110    EAGER_NO_GRAD_DEPTH.with(|depth| depth.get() == 0)
111}
112
113pub(crate) fn eager_capture_active() -> bool {
114    EAGER_CAPTURE_DEPTH.with(|depth| depth.get() > 0)
115}
116
117fn eager_semantic_vjp_enabled() -> bool {
118    #[cfg(test)]
119    if let Some(value) = EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
120        return value;
121    }
122
123    // Semantic eager VJP/JVP on by default (Unification 7).
124    // Set TENFERRO_EAGER_SEMANTIC_VJP=0 to disable.
125    static ENABLED: OnceLock<bool> = OnceLock::new();
126    *ENABLED.get_or_init(|| env::var("TENFERRO_EAGER_SEMANTIC_VJP").map_or(true, |v| v != "0"))
127}
128
129/// Scope guard that temporarily disables eager operation recording.
130///
131/// Values computed while this guard is alive are concrete eager tensors, but
132/// they do not participate in reverse-mode gradient tracking.
133///
134/// # Examples
135///
136/// ```
137/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
138/// use tenferro_cpu::CpuBackend;
139///
140/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
141/// let x = EagerTensor::requires_grad_in(
142///     Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
143///     ctx.clone(),
144/// )?;
145/// let y = {
146///     let _guard = ctx.no_grad();
147///     x.mul(&x)?
148/// };
149/// assert!(!y.tracks_grad());
150/// # Ok::<(), tenferro_ad::Error>(())
151/// ```
152#[derive(Debug)]
153pub struct EagerNoGradGuard {
154    active: bool,
155    // Thread-local depth guard: must not be Send so it cannot be moved to and
156    // dropped on another thread (which would corrupt the creator's depth).
157    _not_send: PhantomData<Rc<()>>,
158}
159
160impl Drop for EagerNoGradGuard {
161    fn drop(&mut self) {
162        if !self.active {
163            return;
164        }
165        EAGER_NO_GRAD_DEPTH.with(|depth| {
166            depth.set(depth.get().saturating_sub(1));
167        });
168        self.active = false;
169    }
170}
171
172/// Scope guard that keeps semantic-trace recording active for untracked
173/// intermediates.
174///
175/// Under active-edge semantics (issue #1665 Def 1), an operation whose inputs
176/// are all untracked produces no autograd nodes and drops its semantic trace.
177/// Inside this guard, such operations still record their semantic trace, so a
178/// later functional JVP/VJP can differentiate with respect to an untracked or
179/// detached leaf. This replaces the pre-Def-1 implicit recording.
180///
181/// # Examples
182///
183/// ```
184/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
185/// use tenferro_cpu::CpuBackend;
186///
187/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
188/// let x = EagerTensor::from_tensor_in(
189///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
190///     ctx.clone(),
191/// )?;
192/// let (y, x) = {
193///     let _capture = ctx.capture_trace();
194///     let y = x.mul(&x)?;
195///     (y, x)
196/// };
197/// let seed = EagerTensor::from_tensor_in(
198///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 1.0]).unwrap(),
199///     ctx.clone(),
200/// )?;
201/// let dx = ctx.vjp(&y, &x, &seed)?;
202/// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[2.0, 4.0]);
203/// # Ok::<(), tenferro_ad::Error>(())
204/// ```
205#[derive(Debug)]
206pub struct EagerTraceCaptureGuard {
207    active: bool,
208    // Thread-local depth guard: must not be Send so it cannot be moved to and
209    // dropped on another thread (which would corrupt the creator's depth).
210    _not_send: PhantomData<Rc<()>>,
211}
212
213impl Drop for EagerTraceCaptureGuard {
214    fn drop(&mut self) {
215        if !self.active {
216            return;
217        }
218        EAGER_CAPTURE_DEPTH.with(|depth| {
219            depth.set(depth.get().saturating_sub(1));
220        });
221        self.active = false;
222    }
223}
224
225pub(crate) fn eager_op_profile_enabled() -> bool {
226    #[cfg(test)]
227    if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
228        return value;
229    }
230
231    static ENABLED: OnceLock<bool> = OnceLock::new();
232    *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
233}
234
235pub(crate) fn eager_op_profile_start() -> Option<Instant> {
236    eager_op_profile_enabled().then(Instant::now)
237}
238
239pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
240    if !eager_op_profile_enabled() {
241        return;
242    }
243    EAGER_OP_PROFILE_STATE.with(|state| {
244        let mut state = state.borrow_mut();
245        let entry = state.entry(section).or_default();
246        entry.calls += 1;
247        entry.total_time += elapsed;
248    });
249}
250
251pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
252    if !eager_op_profile_enabled() {
253        return f();
254    }
255    let started = Instant::now();
256    let result = f();
257    record_eager_op_profile(section, started.elapsed());
258    result
259}
260
261pub(crate) fn maybe_print_eager_op_profile() {
262    if !eager_op_profile_enabled() {
263        return;
264    }
265    let Some(print_every) = eager_op_profile_print_every() else {
266        return;
267    };
268    if print_every == 0 {
269        return;
270    }
271
272    let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
273        state
274            .borrow()
275            .get("nary_op.total")
276            .is_some_and(|entry| entry.calls % print_every == 0)
277    });
278    if should_print {
279        print_and_reset_eager_op_profile();
280    }
281}
282
283fn eager_op_profile_print_every() -> Option<usize> {
284    #[cfg(test)]
285    if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
286        return value;
287    }
288
289    env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
290        .ok()?
291        .parse()
292        .ok()
293}
294
295pub(crate) fn print_and_reset_eager_op_profile() {
296    EAGER_OP_PROFILE_STATE.with(|state| {
297        let mut entries: Vec<_> = state
298            .borrow()
299            .iter()
300            .map(|(section, entry)| (*section, entry.clone()))
301            .collect();
302        state.borrow_mut().clear();
303        entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
304
305        eprintln!("=== tenferro eager op profile ===");
306        for (section, entry) in entries {
307            let Some(per_call_us) = eager_op_profile_per_call_us(&entry) else {
308                continue;
309            };
310            eprintln!(
311                "{section}: calls={} total={:.6}ms per_call={:.3}us",
312                entry.calls,
313                entry.total_time.as_secs_f64() * 1.0e3,
314                per_call_us,
315            );
316        }
317    });
318}
319
320fn eager_op_profile_per_call_us(entry: &EagerOpProfileEntry) -> Option<f64> {
321    (entry.calls != 0).then(|| entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64)
322}
323
324fn runtime_config_error(op: &'static str, source: RuntimeConfigError) -> Error {
325    Error::runtime_state_source(op, ErrorPhase::Execution, source)
326}
327
328fn runtime_state_source<E>(op: &'static str, source: E) -> Error
329where
330    E: std::error::Error + Send + Sync + 'static,
331{
332    Error::runtime_state_source(op, ErrorPhase::Execution, source)
333}
334
335fn cpu_runtime_bridge_unsupported(message: impl Into<String>) -> Error {
336    Error::unsupported(
337        "CpuPlacementBoundEager::refresh_runtime_selection",
338        ErrorPhase::Execution,
339        message,
340    )
341}
342
343fn select_cpu_runtime(runtime: &Runtime) -> Result<CpuRuntimeSelection> {
344    let snapshot = runtime
345        .snapshot()
346        .map_err(|source| runtime_state_source("EagerRuntime::runtime_snapshot", source))?;
347    let engine_id = cpu_runtime_engine_id()
348        .map_err(|source| runtime_config_error("EagerRuntime::cpu_runtime_engine_id", source))?;
349    let expected_hardware = cpu_runtime_hardware_class().map_err(|source| {
350        runtime_config_error("EagerRuntime::cpu_runtime_hardware_class", source)
351    })?;
352    let engine = snapshot
353        .engine(&engine_id)
354        .ok_or_else(|| cpu_runtime_bridge_unsupported("missing CPU runtime engine"))?;
355    validate_cpu_runtime_engine(
356        engine.context_identity(),
357        engine.hardware_class(),
358        engine.capabilities(),
359        &expected_hardware,
360    )?;
361    let epoch = snapshot.epoch();
362    let registration_identity = engine.registration_identity();
363    let capabilities = engine.capabilities().clone();
364    Ok(CpuRuntimeSelection {
365        snapshot,
366        epoch,
367        engine_id,
368        registration_identity,
369        capabilities,
370    })
371}
372
373fn validate_cpu_runtime_engine(
374    context_identity: ExecutionContextIdentity,
375    hardware_class: &HardwareClassId,
376    capabilities: &CoreCapabilityBundle,
377    expected_hardware: &HardwareClassId,
378) -> Result<()> {
379    if context_identity != ExecutionContextIdentity::of::<CpuBackend>() {
380        return Err(cpu_runtime_bridge_unsupported(
381            "CPU runtime context mismatch",
382        ));
383    }
384    if hardware_class != expected_hardware {
385        return Err(cpu_runtime_bridge_unsupported(
386            "CPU runtime hardware mismatch",
387        ));
388    }
389    if capabilities.elementwise().is_none() {
390        return Err(cpu_runtime_bridge_unsupported(
391            "missing CPU runtime capability: elementwise",
392        ));
393    }
394    if capabilities.reduction().is_none() {
395        return Err(cpu_runtime_bridge_unsupported(
396            "missing CPU runtime capability: reduction",
397        ));
398    }
399    if capabilities.indexing().is_none() {
400        return Err(cpu_runtime_bridge_unsupported(
401            "missing CPU runtime capability: indexing",
402        ));
403    }
404    if capabilities.dot_general().is_none() {
405        return Err(cpu_runtime_bridge_unsupported(
406            "missing CPU runtime capability: dot_general",
407        ));
408    }
409    if capabilities.layout().is_none() {
410        return Err(cpu_runtime_bridge_unsupported(
411            "missing CPU runtime capability: layout",
412        ));
413    }
414    Ok(())
415}
416
417/// Stats for caches owned by an [`EagerRuntime`].
418///
419/// `retained_bytes` fields are logical payload estimates, not process RSS.
420#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
421pub struct EagerRuntimeCacheStats {
422    /// Generic extension runtime caches.
423    pub extensions: CacheStats,
424    /// Eager AD transform memoization cache.
425    pub ad_transforms: CacheStats,
426    /// Prepared eager derivative program cache.
427    pub prepared_derivatives: CacheStats,
428}
429
430#[cfg(test)]
431pub(crate) struct EagerGraphExecution {
432    pub(crate) outputs: Vec<Tensor>,
433}
434
435/// A read-only value view retained by an eager tensor record.
436///
437/// The guard borrows the record's allocation group. It never owns a tensor and
438/// cannot be converted into a mutable view.
439///
440/// # Examples
441///
442/// ```
443/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
444/// use tenferro_cpu::CpuBackend;
445///
446/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
447/// let value = EagerTensor::from_tensor_in(
448///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?,
449///     ctx,
450/// )?;
451/// let view = value.value()?;
452/// assert_eq!(view.shape(), &[2]);
453/// # Ok::<(), tenferro_ad::Error>(())
454/// ```
455#[derive(Debug)]
456pub struct ValueGuard<'a> {
457    view: TensorView<'a>,
458}
459
460impl<'a> ValueGuard<'a> {
461    /// Return the scalar dtype of the retained value.
462    pub fn dtype(&self) -> DType {
463        self.view.dtype()
464    }
465
466    /// Return the logical shape of the retained value.
467    pub fn shape(&self) -> &[usize] {
468        self.view.shape()
469    }
470
471    /// Borrow the dtype-erased tensor view.
472    pub fn as_tensor_view(&self) -> &TensorView<'_> {
473        &self.view
474    }
475
476    /// Borrow compact host bytes through the tensor's explicit scalar type.
477    ///
478    /// Backend-resident values return the backend's typed host-access error;
479    /// this method does not download storage implicitly.
480    ///
481    /// # Errors
482    ///
483    /// Returns [`tenferro_tensor::ValidationError::DTypeMismatch`] when
484    /// `T` does not match the view dtype, [`tenferro_tensor::ValidationError::NonContiguousViewAsSlice`]
485    /// for a non-contiguous view, or [`tenferro_tensor::Error::HostAccess`]
486    /// when backend storage cannot be mapped as a host slice.
487    pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&'a [T]> {
488        self.view.as_slice()
489    }
490
491    fn duplicate_host_tensor(&self) -> tenferro_tensor::Result<Tensor> {
492        match &self.view {
493            TensorView::F32(view) => {
494                <f32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
495            }
496            TensorView::F64(view) => {
497                <f64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
498            }
499            TensorView::I32(view) => {
500                <i32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
501            }
502            TensorView::I64(view) => {
503                <i64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
504            }
505            TensorView::Bool(view) => <bool as TensorScalar>::into_tensor(
506                view.shape().to_vec(),
507                view.as_slice()?.to_vec(),
508            ),
509            TensorView::C32(view) => <num_complex::Complex32 as TensorScalar>::into_tensor(
510                view.shape().to_vec(),
511                view.as_slice()?.to_vec(),
512            ),
513            TensorView::C64(view) => <num_complex::Complex64 as TensorScalar>::into_tensor(
514                view.shape().to_vec(),
515                view.as_slice()?.to_vec(),
516            ),
517        }
518    }
519}
520
521/// Read-only retained gradient value.
522///
523/// # Examples
524///
525/// ```
526/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
527/// use tenferro_cpu::CpuBackend;
528///
529/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
530/// let x = EagerTensor::requires_grad_in(
531///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?,
532///     ctx,
533/// )?;
534/// let loss = x.mul(&x)?.reduce_sum(Some(&[0]))?;
535/// let _gradients = loss.backward()?;
536/// let gradient = x.grad()?.expect("tracked leaf has a gradient");
537/// assert_eq!(gradient.shape(), &[2]);
538/// # Ok::<(), tenferro_ad::Error>(())
539/// ```
540#[derive(Clone, Debug)]
541pub struct GradientValue {
542    record: Arc<AdValueRecord>,
543    ctx: Arc<EagerRuntime>,
544}
545
546impl GradientValue {
547    /// Return the scalar dtype of the gradient.
548    pub fn dtype(&self) -> DType {
549        self.record.dtype()
550    }
551
552    /// Return the logical shape of the gradient.
553    pub fn shape(&self) -> &[usize] {
554        self.record.shape()
555    }
556
557    /// Borrow the gradient's value guard.
558    ///
559    /// # Errors
560    ///
561    /// Returns [`Error::RuntimeState`] when the retained gradient record is
562    /// unavailable or its allocation-group descriptor is invalid.
563    pub fn value(&self) -> Result<ValueGuard<'_>> {
564        self.record.value("GradientValue::value")
565    }
566
567    /// Borrow the gradient as a dtype-erased read target.
568    ///
569    /// # Errors
570    ///
571    /// Returns [`Error::RuntimeState`] when the retained gradient record or
572    /// its allocation-group descriptor is unavailable.
573    pub fn tensor_read(&self) -> Result<TensorRead<'_>> {
574        self.record.tensor_read("GradientValue::tensor_read")
575    }
576
577    /// Borrow a compact host slice without downloading backend storage.
578    ///
579    /// # Errors
580    ///
581    /// Returns [`Error::RuntimeState`] when the retained value is unavailable,
582    /// [`tenferro_tensor::ValidationError::DTypeMismatch`] when `T` does
583    /// not match the gradient dtype, or [`tenferro_tensor::Error::HostAccess`]
584    /// when backend storage cannot be mapped as a host slice.
585    pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&[T]> {
586        self.record
587            .value("GradientValue::as_slice")
588            .map_err(|error| {
589                tenferro_tensor::Error::runtime_state_source("GradientValue::as_slice", error)
590            })?
591            .as_slice()
592    }
593
594    /// Explicitly copy a host-resident gradient into a standalone tensor.
595    ///
596    /// # Errors
597    ///
598    /// Returns [`Error::RuntimeState`] when the retained value or execution
599    /// session is unavailable, or a typed backend/host-access error when the
600    /// gradient cannot be materialized as a contiguous tensor.
601    pub fn to_tensor(&self) -> Result<Tensor> {
602        let value = self
603            .record
604            .value("GradientValue::to_tensor")
605            .map_err(|error| {
606                Error::runtime_state_source(
607                    "GradientValue::to_tensor",
608                    ErrorPhase::Execution,
609                    error,
610                )
611            })?;
612        match value.duplicate_host_tensor() {
613            Ok(tensor) => Ok(tensor),
614            Err(_) => {
615                let read = self.record.tensor_read("GradientValue::to_tensor")?;
616                self.ctx
617                    .with_execution_session(|session| session.to_contiguous_read(read))?
618                    .map_err(Error::from)
619            }
620        }
621    }
622}
623
624/// Move-only accumulated gradient bundle backed by one allocation group.
625///
626/// # Examples
627///
628/// ```
629/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
630/// use tenferro_cpu::CpuBackend;
631///
632/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
633/// let x = EagerTensor::requires_grad_in(
634///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?,
635///     ctx,
636/// )?;
637/// let loss = x.mul(&x)?.reduce_sum(Some(&[0]))?;
638/// let gradients = loss.backward()?;
639/// assert!(!gradients.is_empty());
640/// # Ok::<(), tenferro_ad::Error>(())
641/// ```
642#[derive(Debug)]
643pub struct Gradients {
644    group: AllocationGroup,
645    slots: HashMap<ValueKey<StdTensorOp>, DescriptorSlot>,
646}
647
648impl Gradients {
649    fn from_tensors(tensors: HashMap<ValueKey<StdTensorOp>, Tensor>) -> Result<Self> {
650        let (keys, values): (Vec<_>, Vec<_>) = tensors.into_iter().unzip();
651        let (group, bindings) = AllocationGroup::from_tensors(values).map_err(|error| {
652            Error::runtime_state_source("Gradients::from_tensors", ErrorPhase::Execution, error)
653        })?;
654        let slots = keys.into_iter().zip(bindings).collect();
655        Ok(Self { group, slots })
656    }
657
658    /// Return the number of retained gradient descriptors.
659    pub fn len(&self) -> usize {
660        self.slots.len()
661    }
662
663    /// Return whether no gradient was produced.
664    pub fn is_empty(&self) -> bool {
665        self.slots.is_empty()
666    }
667
668    /// Borrow one gradient view by its local value key.
669    pub fn grad(&self, key: &ValueKey<StdTensorOp>) -> Option<TensorView<'_>> {
670        let slot = self.slots.get(key).copied()?;
671        let mut reads = self.group.read_views(std::slice::from_ref(&slot)).ok()?;
672        match reads.pop()? {
673            TensorRead::View(view) => Some(view),
674            TensorRead::Tensor(_) => None,
675        }
676    }
677
678    /// Consume one gradient owner while leaving the bundle unchanged on failure.
679    ///
680    /// # Errors
681    ///
682    /// Returns [`tenferro_tensor::Error::RuntimeState`] when the descriptor is
683    /// invalid or its allocation is aliased. A missing key is reported as
684    /// `Ok(None)`.
685    pub fn take_grad(
686        &mut self,
687        key: &ValueKey<StdTensorOp>,
688    ) -> tenferro_tensor::Result<Option<Tensor>> {
689        let Some(&slot) = self.slots.get(key) else {
690            return Ok(None);
691        };
692        let tensor = self.group.take_tensor(slot).map_err(|error| {
693            tenferro_tensor::Error::runtime_state_source("Gradients::take_grad", error)
694        })?;
695        self.slots.remove(key);
696        Ok(Some(tensor))
697    }
698}
699
700/// Error returned when a value cannot be consumed without changing its owner.
701///
702/// # Examples
703///
704/// ```
705/// use tenferro_ad::{EagerRuntime, EagerTensor, IntoValueError, Tensor};
706/// use tenferro_cpu::CpuBackend;
707///
708/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
709/// let value = EagerTensor::from_tensor_in(
710///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?,
711///     ctx,
712/// )?;
713/// let _shared = value.clone();
714/// assert!(matches!(
715///     value.into_value(),
716///     Err(IntoValueError::NotUnique(_))
717/// ));
718/// # Ok::<(), tenferro_ad::Error>(())
719/// ```
720#[derive(Debug)]
721pub enum IntoValueError<H> {
722    /// Another eager handle, tape record, or checkpoint retains the value.
723    NotUnique(H),
724    /// Group extraction failed after the handle was uniquely acquired.
725    Extract { value: H, error: GroupError },
726}
727
728impl<H> std::fmt::Display for IntoValueError<H> {
729    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
730        match self {
731            Self::NotUnique(_) => formatter.write_str("eager value is retained by another handle"),
732            Self::Extract { error, .. } => {
733                write!(formatter, "eager value extraction failed: {error}")
734            }
735        }
736    }
737}
738
739impl<H: std::fmt::Debug + Send + Sync + 'static> std::error::Error for IntoValueError<H> {}
740
741/// One direct retention container owns the physical allocation group.
742#[derive(Debug)]
743struct RetentionContainer {
744    group: AllocationGroup,
745}
746
747/// Read-only descriptor record used by eager handles and the AD registries.
748#[derive(Debug)]
749pub(crate) struct AdValueRecord {
750    container: Arc<RetentionContainer>,
751    slot: DescriptorSlot,
752    dtype: DType,
753    shape: Box<[usize]>,
754}
755
756impl AdValueRecord {
757    fn from_group(
758        group: AllocationGroup,
759        slot: DescriptorSlot,
760        dtype: DType,
761        shape: Vec<usize>,
762    ) -> Arc<Self> {
763        Arc::new(Self {
764            container: Arc::new(RetentionContainer { group }),
765            slot,
766            dtype,
767            shape: shape.into_boxed_slice(),
768        })
769    }
770
771    fn from_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
772        let dtype = tensor.dtype();
773        let shape = tensor.shape().to_vec();
774        let (group, bindings) = AllocationGroup::from_tensors(vec![tensor])
775            .map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
776        let slot = bindings.first().copied().ok_or_else(|| {
777            Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
778        })?;
779        Ok(Self::from_group(group, slot, dtype, shape))
780    }
781
782    fn tensor_read(&self, op: &'static str) -> Result<TensorRead<'_>> {
783        let mut reads = self
784            .container
785            .group
786            .read_views(std::slice::from_ref(&self.slot))
787            .map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
788        reads.pop().ok_or_else(|| {
789            Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
790        })
791    }
792
793    fn value(&self, op: &'static str) -> Result<ValueGuard<'_>> {
794        match self.tensor_read(op)? {
795            TensorRead::View(view) => Ok(ValueGuard { view }),
796            TensorRead::Tensor(_) => Err(Error::runtime_state(
797                op,
798                ErrorPhase::Execution,
799                "allocation-group value did not produce a borrowed descriptor view",
800            )),
801        }
802    }
803
804    fn dtype(&self) -> DType {
805        self.dtype
806    }
807
808    fn shape(&self) -> &[usize] {
809        &self.shape
810    }
811}
812
813/// Placement-selected CPU view of one [`EagerRuntime`].
814///
815/// The view snapshots the runtime's CPU coordinator/provider bundle and the
816/// immutable runtime registration metadata when [`EagerRuntime::on_cpu`] is
817/// called. It holds no resource permit while idle and enters one backend
818/// session only while [`Self::with_eager_session`] runs. The session exposes
819/// core [`BackendSession`] operations on concrete [`Tensor`] values. This
820/// bridge deliberately does not expose the eager runtime's linalg, FFT, einsum,
821/// or extension-runtime registries.
822///
823/// The value is intentionally not `Clone`: mutable use makes concurrent
824/// session ownership explicit without adding another backend mutex.
825///
826/// # Examples
827///
828/// ```rust
829/// use tenferro_ad::EagerRuntime;
830/// use tenferro_cpu::CpuPlacement;
831///
832/// let runtime = EagerRuntime::new()?;
833/// let cpu = runtime.on_cpu(CpuPlacement::Auto)?;
834/// assert_eq!(cpu.runtime_id(), runtime.id());
835/// # Ok::<(), tenferro_ad::Error>(())
836/// ```
837pub struct CpuPlacementBoundEager {
838    runtime: Arc<EagerRuntime>,
839    backend: CpuBackend,
840    snapshot: Arc<RuntimeConfigSnapshot>,
841    epoch: RuntimeEpoch,
842    engine_id: EngineId,
843    registration_identity: RegistrationIdentity,
844    capabilities: CoreCapabilityBundle,
845}
846
847impl fmt::Debug for CpuPlacementBoundEager {
848    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
849        f.debug_struct("CpuPlacementBoundEager")
850            .field("runtime_id", &self.runtime.id())
851            .field("placement", &self.backend.placement())
852            .field("runtime_epoch", &self.epoch)
853            .field("engine_id", &self.engine_id)
854            .field("registration_identity", &self.registration_identity)
855            .finish_non_exhaustive()
856    }
857}
858
859impl CpuPlacementBoundEager {
860    fn refresh_runtime_selection(&mut self) -> Result<()> {
861        let current_epoch = self.runtime.runtime.epoch().map_err(|source| {
862            runtime_state_source("CpuPlacementBoundEager::refresh_runtime_selection", source)
863        })?;
864        if current_epoch == self.epoch {
865            return Ok(());
866        }
867
868        #[cfg(test)]
869        CPU_RUNTIME_SELECTION_REFRESHES.fetch_add(1, Ordering::SeqCst);
870
871        let selection = select_cpu_runtime(&self.runtime.runtime)?;
872        self.snapshot = selection.snapshot;
873        self.epoch = selection.epoch;
874        self.engine_id = selection.engine_id;
875        self.registration_identity = selection.registration_identity;
876        self.capabilities = selection.capabilities;
877        Ok(())
878    }
879
880    /// Return the identity of the original eager runtime.
881    ///
882    /// # Examples
883    ///
884    /// ```rust
885    /// use tenferro_ad::EagerRuntime;
886    /// use tenferro_cpu::CpuPlacement;
887    ///
888    /// let runtime = EagerRuntime::new()?;
889    /// let cpu = runtime.on_cpu(CpuPlacement::Auto)?;
890    /// assert_eq!(cpu.runtime_id(), runtime.id());
891    /// # Ok::<(), tenferro_ad::Error>(())
892    /// ```
893    pub fn runtime_id(&self) -> ContextId {
894        self.runtime.id()
895    }
896
897    /// Return the placement requested when this view was created.
898    ///
899    /// # Examples
900    ///
901    /// ```rust
902    /// use tenferro_ad::EagerRuntime;
903    /// use tenferro_cpu::CpuPlacement;
904    ///
905    /// let runtime = EagerRuntime::new()?;
906    /// let cpu = runtime.on_cpu(CpuPlacement::Auto)?;
907    /// assert_eq!(cpu.placement(), CpuPlacement::Auto);
908    /// # Ok::<(), tenferro_ad::Error>(())
909    /// ```
910    pub fn placement(&self) -> CpuPlacement {
911        self.backend.placement()
912    }
913
914    /// Enter one CPU backend session and run core operations through it.
915    ///
916    /// One call creates one backend session. Tenferro-managed CPU executors
917    /// enter once around the closure and core operations reuse that compatible
918    /// execution scope. The closure may borrow stack data and need not be
919    /// `'static`.
920    ///
921    /// This phase-2 bridge accepts only core [`BackendSession`] operations. It
922    /// does not lock or dispatch the eager runtime's linalg, FFT, einsum, or
923    /// extension registries.
924    ///
925    /// # Examples
926    ///
927    /// ```rust
928    /// use tenferro_ad::{EagerRuntime, Error};
929    /// use tenferro_cpu::CpuPlacement;
930    /// use tenferro_tensor::{Tensor, TensorElementwise};
931    ///
932    /// let runtime = EagerRuntime::new()?;
933    /// let mut cpu = runtime.on_cpu(CpuPlacement::Auto)?;
934    /// let lhs = Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?;
935    /// let rhs = Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?;
936    /// let output = cpu.with_eager_session(|session| {
937    ///     TensorElementwise::add(session, &lhs, &rhs).map_err(Error::from)
938    /// })?;
939    /// assert_eq!(output.as_slice::<f64>().unwrap(), &[3.0]);
940    /// # Ok::<(), Error>(())
941    /// ```
942    ///
943    /// # Errors
944    ///
945    /// Returns the callback's [`Error`] unchanged. Core backend operations may
946    /// report validation, unsupported capability, backend, or runtime-state
947    /// failures through that error.
948    ///
949    /// # Panics
950    ///
951    /// The existing CPU backend re-entry guard panics if the callback enters a
952    /// public `CpuBackend` or calls an ordinary `EagerTensor` operation on this
953    /// same runtime. Use only the borrowed `session` for work inside the scope.
954    pub fn with_eager_session<R: Send>(
955        &mut self,
956        f: impl FnOnce(&mut dyn BackendSession) -> Result<R> + Send,
957    ) -> Result<R> {
958        self.refresh_runtime_selection()?;
959        self.backend.with_backend_session(f)
960    }
961}
962
963/// Shared eager execution context for tensors on a backend.
964///
965/// Reusing one context lets eager tensors share backend state, extension
966/// runtime caches, and gradient storage across a computation.
967///
968/// # Examples
969///
970/// ```
971/// use tenferro_cpu::CpuBackend;
972/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
973///
974/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
975/// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
976/// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx).unwrap();
977/// let z = x.add(&y).unwrap();
978///
979/// assert_eq!(z.value().unwrap().as_slice::<f64>().unwrap(), &[3.0]);
980/// # Ok::<(), tenferro_ad::Error>(())
981/// ```
982pub struct EagerRuntime {
983    id: ContextId,
984    runtime: Runtime,
985    // The backend and its exact runtime engine registration are selected
986    // together during construction and remain paired for this runtime's
987    // lifetime. The mutex only serializes mutable backend operations.
988    backend: Mutex<EagerBackend>,
989    extension_install_lock: Mutex<()>,
990    pub(crate) extension_caches: Mutex<ExtensionCacheStore>,
991    semantic_extension_rules: SemanticExtensionRuleSet,
992    grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
993    value_records: Mutex<HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>,
994    ad_transform_cache: Arc<AdTransformCache>,
995    /// S2: prepared derivative programs keyed by semantic structure, wrt input,
996    /// and concrete bound input metadata. Avoids re-running freeze+AD
997    /// transform+compile_frozen on warm structure hits.
998    prepared_derivative_cache: Mutex<PreparedDerivativeCache>,
999}
1000
1001impl fmt::Debug for EagerRuntime {
1002    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1003        let mut debug = f.debug_struct("EagerRuntime");
1004        debug.field("id", &self.id);
1005        debug.field("runtime_id", &self.runtime.id());
1006        debug.field("runtime_epoch", &self.runtime.epoch().ok());
1007        match self.backend.try_lock() {
1008            Ok(backend) => {
1009                debug.field("backend", &*backend);
1010            }
1011            Err(_) => {
1012                debug.field("backend", &"<locked>");
1013            }
1014        }
1015        match self.extension_caches.try_lock() {
1016            Ok(caches) => {
1017                debug.field(
1018                    "extension_cache_stats",
1019                    &caches.stats(ExtensionCacheSelector::All),
1020                );
1021            }
1022            Err(_) => {
1023                debug.field("extension_cache_stats", &"<locked>");
1024            }
1025        }
1026        match self.extension_install_lock.try_lock() {
1027            Ok(_) => {
1028                debug.field("extension_install_lock", &"<unlocked>");
1029            }
1030            Err(_) => {
1031                debug.field("extension_install_lock", &"<locked>");
1032            }
1033        }
1034        debug.field("semantic_extension_rules", &self.semantic_extension_rules);
1035        match self.grad_slots.try_lock() {
1036            Ok(slots) => {
1037                debug.field("grad_slots_len", &slots.len());
1038            }
1039            Err(_) => {
1040                debug.field("grad_slots_len", &"<locked>");
1041            }
1042        }
1043        match self.value_records.try_lock() {
1044            Ok(records) => {
1045                debug.field("value_records_len", &records.len());
1046            }
1047            Err(_) => {
1048                debug.field("value_records_len", &"<locked>");
1049            }
1050        }
1051        match self.ad_transform_cache.stats() {
1052            Ok(stats) => {
1053                debug.field("ad_transform_cache_stats", &stats);
1054            }
1055            Err(err) => {
1056                debug.field("ad_transform_cache_stats", &format_args!("{err}"));
1057            }
1058        }
1059        match self.prepared_derivative_cache.try_lock() {
1060            Ok(cache) => {
1061                debug.field("prepared_derivative_cache_stats", &cache.stats());
1062            }
1063            Err(_) => {
1064                debug.field("prepared_derivative_cache_stats", &"<locked>");
1065            }
1066        }
1067        debug.finish_non_exhaustive()
1068    }
1069}
1070
1071impl EagerRuntime {
1072    pub(crate) fn lock_backend(&self) -> Result<MutexGuard<'_, EagerBackend>> {
1073        self.backend.lock().map_err(|_| {
1074            Error::runtime_state("eager_backend", ErrorPhase::Execution, "lock poisoned")
1075        })
1076    }
1077
1078    fn lock_extension_caches(&self) -> Result<MutexGuard<'_, ExtensionCacheStore>> {
1079        self.extension_caches.lock().map_err(|_| {
1080            Error::runtime_state(
1081                "eager_extension_caches",
1082                ErrorPhase::Execution,
1083                "lock poisoned",
1084            )
1085        })
1086    }
1087
1088    fn lock_extension_install(&self) -> Result<MutexGuard<'_, ()>> {
1089        self.extension_install_lock.lock().map_err(|_| {
1090            Error::runtime_state(
1091                "eager_extension_install",
1092                ErrorPhase::Execution,
1093                "lock poisoned",
1094            )
1095        })
1096    }
1097
1098    fn lock_prepared_derivative_cache(&self) -> Result<MutexGuard<'_, PreparedDerivativeCache>> {
1099        self.prepared_derivative_cache.lock().map_err(|_| {
1100            Error::runtime_state(
1101                "prepared_derivative_cache",
1102                ErrorPhase::Execution,
1103                "lock poisoned",
1104            )
1105        })
1106    }
1107
1108    fn lock_grad_slots(
1109        &self,
1110    ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
1111        self.grad_slots.lock().map_err(|_| {
1112            Error::runtime_state(
1113                "eager_gradient_slots",
1114                ErrorPhase::Execution,
1115                "lock poisoned",
1116            )
1117        })
1118    }
1119
1120    fn lock_value_records(
1121        &self,
1122    ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>> {
1123        self.value_records.lock().map_err(|_| {
1124            Error::runtime_state(
1125                "eager_value_registry",
1126                ErrorPhase::Execution,
1127                "lock poisoned",
1128            )
1129        })
1130    }
1131
1132    fn from_backend(backend: EagerBackend) -> Result<Self> {
1133        Self::from_backend_with_rules_and_cache(
1134            backend,
1135            SemanticExtensionRuleSet::default(),
1136            Arc::new(AdTransformCache::new()),
1137        )
1138    }
1139
1140    fn from_backend_with_rules_and_cache(
1141        backend: EagerBackend,
1142        semantic_extension_rules: SemanticExtensionRuleSet,
1143        ad_transform_cache: Arc<AdTransformCache>,
1144    ) -> Result<Self> {
1145        let runtime = eager_runtime_for_backend(&backend)
1146            .map_err(|source| runtime_config_error("EagerRuntime::from_backend", source))?;
1147        Ok(Self {
1148            id: ContextId::fresh(),
1149            runtime,
1150            backend: Mutex::new(backend),
1151            extension_install_lock: Mutex::new(()),
1152            extension_caches: Mutex::new(ExtensionCacheStore::new()),
1153            semantic_extension_rules,
1154            grad_slots: Mutex::new(HashMap::new()),
1155            value_records: Mutex::new(HashMap::new()),
1156            ad_transform_cache,
1157            prepared_derivative_cache: Mutex::new(PreparedDerivativeCache::default()),
1158        })
1159    }
1160
1161    /// Create a shared CPU eager execution context.
1162    ///
1163    /// # Examples
1164    ///
1165    /// ```
1166    /// use tenferro_ad::EagerRuntime;
1167    ///
1168    /// let ctx = EagerRuntime::new()?;
1169    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
1170    /// # Ok::<(), tenferro_ad::Error>(())
1171    /// ```
1172    ///
1173    /// # Errors
1174    ///
1175    /// Returns [`Error::RuntimeStateSource`] when provider runtime
1176    /// registration cannot be configured, preserving the underlying
1177    /// [`RuntimeConfigError`] as the typed error source.
1178    pub fn new() -> Result<Arc<Self>> {
1179        Self::with_cpu_backend(CpuBackend::new())
1180    }
1181
1182    /// Create a shared eager execution context from a configured CPU backend.
1183    ///
1184    /// # Examples
1185    ///
1186    /// ```
1187    /// use tenferro_cpu::CpuBackend;
1188    /// use tenferro_ad::{EagerRuntime};
1189    ///
1190    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::with_threads(1)?)?;
1191    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
1192    /// # Ok::<(), Box<dyn std::error::Error>>(())
1193    /// ```
1194    ///
1195    /// # Errors
1196    ///
1197    /// Returns [`Error::RuntimeStateSource`] when provider runtime
1198    /// registration cannot be configured, preserving the underlying
1199    /// [`RuntimeConfigError`] as the typed error source.
1200    pub fn with_cpu_backend(backend: CpuBackend) -> Result<Arc<Self>> {
1201        Ok(Arc::new(Self::from_backend(EagerBackend::cpu(backend))?))
1202    }
1203
1204    /// Snapshot a placement-selected CPU handle from this eager runtime.
1205    ///
1206    /// The eager backend lock is held only long enough to verify the backend
1207    /// kind and clone its CPU coordinator/provider snapshot. Placement
1208    /// resolution happens after that guard is dropped. The returned value does
1209    /// not hold a resource permit or a second runtime/backend mutex while idle.
1210    ///
1211    /// # Examples
1212    ///
1213    /// ```rust
1214    /// use tenferro_ad::EagerRuntime;
1215    /// use tenferro_cpu::CpuPlacement;
1216    ///
1217    /// let runtime = EagerRuntime::new()?;
1218    /// let cpu = runtime.on_cpu(CpuPlacement::Auto)?;
1219    /// assert_eq!(cpu.runtime_id(), runtime.id());
1220    /// # Ok::<(), tenferro_ad::Error>(())
1221    /// ```
1222    ///
1223    /// # Errors
1224    ///
1225    /// Returns [`Error::RuntimeState`] if the eager backend lock is poisoned,
1226    /// [`Error::Unsupported`] if the runtime is not CPU-backed, or a typed
1227    /// tensor runtime error retaining [`tenferro_cpu::CpuPlacementError`] when
1228    /// the requested placement cannot be resolved.
1229    pub fn on_cpu(self: &Arc<Self>, placement: CpuPlacement) -> Result<CpuPlacementBoundEager> {
1230        let backend = {
1231            let backend = self.lock_backend()?;
1232            backend.cpu_snapshot().ok_or_else(|| {
1233                Error::unsupported(
1234                    "EagerRuntime::on_cpu",
1235                    ErrorPhase::Execution,
1236                    "the eager runtime is not CPU-backed",
1237                )
1238            })?
1239        };
1240        let selection = select_cpu_runtime(&self.runtime)?;
1241        let backend = backend.for_placement(placement).map_err(|source| {
1242            let error: tenferro_tensor::Error = CpuBackendError::Placement {
1243                op: "EagerRuntime::on_cpu",
1244                source,
1245            }
1246            .into();
1247            Error::from(error)
1248        })?;
1249        Ok(CpuPlacementBoundEager {
1250            runtime: Arc::clone(self),
1251            backend,
1252            snapshot: selection.snapshot,
1253            epoch: selection.epoch,
1254            engine_id: selection.engine_id,
1255            registration_identity: selection.registration_identity,
1256            capabilities: selection.capabilities,
1257        })
1258    }
1259
1260    /// Create a shared CPU eager context with explicit AD extension rules.
1261    ///
1262    /// # Examples
1263    ///
1264    /// ```rust
1265    /// use tenferro_cpu::CpuBackend;
1266    /// use tenferro_ad::{AdContext, EagerRuntime};
1267    ///
1268    /// let ad = AdContext::builder().build().unwrap();
1269    /// let ctx = EagerRuntime::with_cpu_backend_and_ad_context(CpuBackend::new(), &ad)?;
1270    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
1271    /// # Ok::<(), tenferro_ad::Error>(())
1272    /// ```
1273    ///
1274    /// # Errors
1275    ///
1276    /// Returns [`Error::RuntimeStateSource`] when provider runtime
1277    /// registration cannot be configured, preserving the underlying
1278    /// [`RuntimeConfigError`] as the typed error source.
1279    pub fn with_cpu_backend_and_ad_context(
1280        backend: CpuBackend,
1281        ad: &AdContext,
1282    ) -> Result<Arc<Self>> {
1283        Ok(Arc::new(Self::from_backend_with_rules_and_cache(
1284            EagerBackend::cpu(backend),
1285            ad.semantic_extension_rules().clone(),
1286            ad.ad_transform_cache(),
1287        )?))
1288    }
1289
1290    /// Create a shared eager execution context from a configured CUDA backend.
1291    ///
1292    /// # Examples
1293    ///
1294    /// ```
1295    /// use tenferro_gpu::cuda::CudaBackend;
1296    /// use tenferro_ad::EagerRuntime;
1297    ///
1298    /// let _ctor: fn(CudaBackend) -> tenferro_ad::Result<std::sync::Arc<EagerRuntime>> =
1299    ///     EagerRuntime::with_cuda_backend;
1300    /// ```
1301    #[cfg(feature = "cuda")]
1302    ///
1303    /// # Errors
1304    ///
1305    /// Returns [`Error::RuntimeStateSource`] when provider runtime
1306    /// registration cannot be configured, preserving the underlying
1307    /// [`RuntimeConfigError`] as the typed error source.
1308    pub fn with_cuda_backend(backend: CudaBackend) -> Result<Arc<Self>> {
1309        Ok(Arc::new(Self::from_backend(EagerBackend::cuda(backend))?))
1310    }
1311
1312    /// Create a shared CUDA eager context with explicit AD extension rules.
1313    ///
1314    /// # Examples
1315    ///
1316    /// ```rust
1317    /// use tenferro_ad::{AdContext, EagerRuntime};
1318    /// use tenferro_gpu::cuda::CudaBackend;
1319    ///
1320    /// let _ctor: fn(CudaBackend, &AdContext) -> tenferro_ad::Result<std::sync::Arc<EagerRuntime>> =
1321    ///     EagerRuntime::with_cuda_backend_and_ad_context;
1322    /// ```
1323    #[cfg(feature = "cuda")]
1324    ///
1325    /// # Errors
1326    ///
1327    /// Returns [`Error::RuntimeStateSource`] when provider runtime
1328    /// registration cannot be configured, preserving the underlying
1329    /// [`RuntimeConfigError`] as the typed error source.
1330    pub fn with_cuda_backend_and_ad_context(
1331        backend: CudaBackend,
1332        ad: &AdContext,
1333    ) -> Result<Arc<Self>> {
1334        Ok(Arc::new(Self::from_backend_with_rules_and_cache(
1335            EagerBackend::cuda(backend),
1336            ad.semantic_extension_rules().clone(),
1337            ad.ad_transform_cache(),
1338        )?))
1339    }
1340
1341    /// Create a shared eager execution context from a configured WebGPU backend.
1342    ///
1343    /// # Examples
1344    ///
1345    /// ```
1346    /// use tenferro_ad::EagerRuntime;
1347    /// use tenferro_gpu::webgpu::WebGpuBackend;
1348    ///
1349    /// let _ctor: fn(WebGpuBackend) -> tenferro_ad::Result<std::sync::Arc<EagerRuntime>> =
1350    ///     EagerRuntime::with_webgpu_backend;
1351    /// ```
1352    #[cfg(feature = "webgpu")]
1353    ///
1354    /// # Errors
1355    ///
1356    /// Returns [`Error::RuntimeStateSource`] when provider runtime
1357    /// registration cannot be configured, preserving the underlying
1358    /// [`RuntimeConfigError`] as the typed error source.
1359    pub fn with_webgpu_backend(backend: WebGpuBackend) -> Result<Arc<Self>> {
1360        Ok(Arc::new(Self::from_backend(EagerBackend::webgpu(backend))?))
1361    }
1362
1363    /// Create a shared WebGPU eager context with explicit AD extension rules.
1364    ///
1365    /// # Examples
1366    ///
1367    /// ```rust
1368    /// use tenferro_ad::{AdContext, EagerRuntime};
1369    /// use tenferro_gpu::webgpu::WebGpuBackend;
1370    ///
1371    /// let _ctor: fn(WebGpuBackend, &AdContext) -> tenferro_ad::Result<std::sync::Arc<EagerRuntime>> =
1372    ///     EagerRuntime::with_webgpu_backend_and_ad_context;
1373    /// ```
1374    #[cfg(feature = "webgpu")]
1375    ///
1376    /// # Errors
1377    ///
1378    /// Returns [`Error::RuntimeStateSource`] when provider runtime
1379    /// registration cannot be configured, preserving the underlying
1380    /// [`RuntimeConfigError`] as the typed error source.
1381    pub fn with_webgpu_backend_and_ad_context(
1382        backend: WebGpuBackend,
1383        ad: &AdContext,
1384    ) -> Result<Arc<Self>> {
1385        Ok(Arc::new(Self::from_backend_with_rules_and_cache(
1386            EagerBackend::webgpu(backend),
1387            ad.semantic_extension_rules().clone(),
1388            ad.ad_transform_cache(),
1389        )?))
1390    }
1391
1392    /// Return an opaque identifier for this context.
1393    ///
1394    /// # Examples
1395    ///
1396    /// ```
1397    /// use tenferro_cpu::CpuBackend;
1398    /// use tenferro_ad::{EagerRuntime};
1399    ///
1400    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1401    /// assert_ne!(ctx.id(), EagerRuntime::with_cpu_backend(CpuBackend::new())?.id());
1402    /// # Ok::<(), tenferro_ad::Error>(())
1403    /// ```
1404    pub fn id(&self) -> ContextId {
1405        self.id
1406    }
1407
1408    /// Disable eager operation recording on the current thread until the guard is dropped.
1409    ///
1410    /// This is useful for optimizer updates, metric calculations, and other
1411    /// eager computations that should not become part of the AD tape.
1412    ///
1413    /// # Examples
1414    ///
1415    /// ```
1416    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1417    /// use tenferro_cpu::CpuBackend;
1418    ///
1419    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1420    /// let x = EagerTensor::requires_grad_in(
1421    ///     Tensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap(),
1422    ///     ctx.clone(),
1423    /// )?;
1424    /// let y = {
1425    ///     let _guard = ctx.no_grad();
1426    ///     x.mul(&x)?
1427    /// };
1428    /// assert!(!y.tracks_grad());
1429    /// # Ok::<(), tenferro_ad::Error>(())
1430    /// ```
1431    pub fn no_grad(&self) -> EagerNoGradGuard {
1432        EAGER_NO_GRAD_DEPTH.with(|depth| {
1433            depth.set(depth.get().saturating_add(1));
1434        });
1435        EagerNoGradGuard {
1436            active: true,
1437            _not_send: PhantomData,
1438        }
1439    }
1440
1441    /// Keep semantic-trace recording active for untracked intermediates.
1442    ///
1443    /// See [`EagerTraceCaptureGuard`] for the full contract and an example.
1444    pub fn capture_trace(&self) -> EagerTraceCaptureGuard {
1445        EAGER_CAPTURE_DEPTH.with(|depth| {
1446            depth.set(depth.get().saturating_add(1));
1447        });
1448        EagerTraceCaptureGuard {
1449            active: true,
1450            _not_send: PhantomData,
1451        }
1452    }
1453
1454    /// Install or replace one extension module on this eager context's runtime.
1455    ///
1456    /// Eager extension wrappers call this as an idempotent "ensure installed"
1457    /// step. When the exact module instance (same module ID and allocation) is
1458    /// already installed, this is a read-only no-op that returns the current
1459    /// runtime epoch without acquiring the install lock or reconfiguring. The
1460    /// cold or replacement paths keep the transactional install-or-replace
1461    /// behavior, serialized so parallel first-use of the same extension family
1462    /// cannot publish over another thread's base snapshot.
1463    ///
1464    /// # Errors
1465    ///
1466    /// Returns [`tenferro_runtime::Error::RuntimeState`] when runtime
1467    /// reconfiguration fails or the extension module transaction is invalid.
1468    pub fn install_extension_module(
1469        &self,
1470        module: Arc<dyn ExtensionModule>,
1471    ) -> Result<RuntimeEpoch> {
1472        let snapshot = self.runtime.snapshot().map_err(|source| {
1473            runtime_state_source("EagerRuntime::install_extension_module", source)
1474        })?;
1475        if snapshot.has_extension_module_identical(&module) {
1476            return Ok(snapshot.epoch());
1477        }
1478        let _install_guard = self.lock_extension_install()?;
1479        self.runtime
1480            .reconfigure(|edit| {
1481                edit.replace_extension_module(module)?;
1482                Ok(())
1483            })
1484            .map_err(|source| {
1485                runtime_state_source("EagerRuntime::install_extension_module", source)
1486            })
1487    }
1488
1489    pub(crate) fn ensure_extension_module_for_engine(
1490        &self,
1491        module: Arc<dyn ExtensionModule>,
1492        family_id: &'static str,
1493        engine_id: &EngineId,
1494    ) -> Result<RuntimeEpoch> {
1495        let snapshot = self.runtime.snapshot().map_err(|source| {
1496            runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
1497        })?;
1498        if snapshot.has_extension_module_engine(module.module_id(), family_id, engine_id) {
1499            return Ok(snapshot.epoch());
1500        }
1501        let _install_guard = self.lock_extension_install()?;
1502        self.runtime
1503            .reconfigure(|edit| {
1504                edit.ensure_extension_module_for_engine(module, family_id, engine_id)?;
1505                Ok(())
1506            })
1507            .map_err(|source| {
1508                runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
1509            })
1510    }
1511
1512    pub(crate) fn runtime(&self) -> &Runtime {
1513        &self.runtime
1514    }
1515
1516    pub(crate) fn eager_extension_target(&self) -> Result<EagerExtensionTarget> {
1517        let (engine_id, backend_kind) = {
1518            let backend = self.lock_backend()?;
1519            match &*backend {
1520                EagerBackend::Cpu(_) => (
1521                    cpu_runtime_engine_id().map_err(|source| {
1522                        runtime_config_error("EagerRuntime::eager_extension_target", source)
1523                    })?,
1524                    EagerExtensionBackendKind::Cpu,
1525                ),
1526                #[cfg(test)]
1527                EagerBackend::Recording(_) => {
1528                    return Err(Error::unsupported(
1529                        "EagerRuntime::eager_extension_target",
1530                        ErrorPhase::Execution,
1531                        "the recording backend has no registered eager extension engine",
1532                    ));
1533                }
1534                #[cfg(feature = "cuda")]
1535                EagerBackend::Cuda(_) => (
1536                    cuda_runtime_engine_id().map_err(|source| {
1537                        runtime_config_error("EagerRuntime::eager_extension_target", source)
1538                    })?,
1539                    EagerExtensionBackendKind::Cuda,
1540                ),
1541                #[cfg(feature = "webgpu")]
1542                EagerBackend::WebGpu(_) => (
1543                    tenferro_gpu::webgpu::webgpu_runtime_engine_id().map_err(|source| {
1544                        runtime_config_error("EagerRuntime::eager_extension_target", source)
1545                    })?,
1546                    EagerExtensionBackendKind::WebGpu,
1547                ),
1548            }
1549        };
1550        let target = EagerExtensionTarget {
1551            engine_id,
1552            backend_kind,
1553        };
1554        validate_eager_extension_target(&self.runtime, &target)?;
1555        Ok(target)
1556    }
1557
1558    /// Clear generic extension runtime cache entries.
1559    ///
1560    /// # Examples
1561    ///
1562    /// ```
1563    /// use tenferro_cpu::CpuBackend;
1564    /// use tenferro_ad::{EagerRuntime};
1565    ///
1566    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1567    /// ctx.clear_extension_caches()?;
1568    /// assert_eq!(ctx.cache_stats()?.extensions.entries, 0);
1569    /// # Ok::<(), tenferro_ad::Error>(())
1570    /// ```
1571    ///
1572    /// # Errors
1573    ///
1574    /// Returns [`tenferro_runtime::Error::RuntimeState`] when the extension
1575    /// cache lock is poisoned.
1576    pub fn clear_extension_caches(&self) -> Result<()> {
1577        self.lock_extension_caches()?.clear();
1578        Ok(())
1579    }
1580
1581    /// Clear every cache owned by this eager context.
1582    ///
1583    /// # Examples
1584    ///
1585    /// ```
1586    /// use tenferro_cpu::CpuBackend;
1587    /// use tenferro_ad::{EagerRuntime};
1588    ///
1589    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1590    /// ctx.clear_caches()?;
1591    /// assert_eq!(ctx.cache_stats()?.extensions.entries, 0);
1592    /// assert_eq!(ctx.cache_stats()?.ad_transforms.entries, 0);
1593    /// assert_eq!(ctx.cache_stats()?.prepared_derivatives.entries, 0);
1594    /// # Ok::<(), tenferro_ad::Error>(())
1595    /// ```
1596    ///
1597    /// # Errors
1598    ///
1599    /// Returns [`tenferro_runtime::Error::RuntimeState`] when either the
1600    /// extension cache or AD-transform cache is poisoned.
1601    pub fn clear_caches(&self) -> Result<()> {
1602        self.clear_extension_caches()?;
1603        self.clear_ad_transform_caches()?;
1604        self.clear_prepared_derivative_cache()?;
1605        Ok(())
1606    }
1607
1608    /// Clear prepared derivative program cache entries.
1609    ///
1610    /// # Examples
1611    ///
1612    /// ```rust
1613    /// use tenferro_ad::EagerRuntime;
1614    /// use tenferro_cpu::CpuBackend;
1615    ///
1616    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1617    /// ctx.clear_prepared_derivative_cache()?;
1618    /// assert_eq!(ctx.cache_stats()?.prepared_derivatives.entries, 0);
1619    /// # Ok::<(), tenferro_ad::Error>(())
1620    /// ```
1621    ///
1622    /// # Errors
1623    ///
1624    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
1625    /// derivative cache lock is poisoned.
1626    pub fn clear_prepared_derivative_cache(&self) -> Result<()> {
1627        self.lock_prepared_derivative_cache()?.clear();
1628        Ok(())
1629    }
1630
1631    /// Return eager runtime cache-entry and retained-byte stats.
1632    ///
1633    /// # Examples
1634    ///
1635    /// ```
1636    /// use tenferro_cpu::CpuBackend;
1637    /// use tenferro_ad::{EagerRuntime};
1638    ///
1639    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1640    /// let stats = ctx.cache_stats()?;
1641    /// assert_eq!(stats.extensions.entries, 0);
1642    /// assert_eq!(stats.ad_transforms.entries, 0);
1643    /// assert_eq!(stats.prepared_derivatives.entries, 0);
1644    /// # Ok::<(), tenferro_ad::Error>(())
1645    /// ```
1646    ///
1647    /// # Errors
1648    ///
1649    /// Returns [`tenferro_runtime::Error::RuntimeState`] when a cache or
1650    /// AD-transform cache lock is poisoned.
1651    pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
1652        Ok(EagerRuntimeCacheStats {
1653            extensions: self
1654                .lock_extension_caches()?
1655                .stats(ExtensionCacheSelector::All),
1656            ad_transforms: self.ad_transform_cache.stats()?,
1657            prepared_derivatives: self.lock_prepared_derivative_cache()?.stats(),
1658        })
1659    }
1660
1661    /// Return the AD transform cache retention limits.
1662    ///
1663    /// # Examples
1664    ///
1665    /// ```
1666    /// use tenferro_ad::EagerRuntime;
1667    /// use tenferro_cpu::CpuBackend;
1668    ///
1669    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1670    /// assert!(ctx.ad_transform_cache_limits()?.max_entries().get() > 0);
1671    /// # Ok::<(), tenferro_ad::Error>(())
1672    /// ```
1673    ///
1674    /// # Errors
1675    ///
1676    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
1677    /// cache lock is poisoned.
1678    pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
1679        self.ad_transform_cache.limits()
1680    }
1681
1682    /// Replace AD transform cache retention limits.
1683    ///
1684    /// # Examples
1685    ///
1686    /// ```
1687    /// use std::num::NonZeroUsize;
1688    /// use tenferro_ad::{AdTransformCacheLimits, EagerRuntime};
1689    /// use tenferro_cpu::CpuBackend;
1690    ///
1691    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1692    /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
1693    /// ctx.set_ad_transform_cache_limits(limits)?;
1694    /// assert_eq!(ctx.ad_transform_cache_limits()?, limits);
1695    /// # Ok::<(), tenferro_ad::Error>(())
1696    /// ```
1697    ///
1698    /// # Errors
1699    ///
1700    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
1701    /// cache lock is poisoned while updating limits.
1702    pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
1703        self.ad_transform_cache.set_limits(limits)
1704    }
1705
1706    /// Clear AD transform cache entries visible through this eager runtime.
1707    ///
1708    /// # Examples
1709    ///
1710    /// ```
1711    /// use tenferro_ad::EagerRuntime;
1712    /// use tenferro_cpu::CpuBackend;
1713    ///
1714    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1715    /// ctx.clear_ad_transform_caches()?;
1716    /// assert_eq!(ctx.cache_stats()?.ad_transforms.entries, 0);
1717    /// # Ok::<(), tenferro_ad::Error>(())
1718    /// ```
1719    ///
1720    /// # Errors
1721    ///
1722    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
1723    /// cache lock is poisoned while clearing entries.
1724    pub fn clear_ad_transform_caches(&self) -> Result<()> {
1725        self.ad_transform_cache.clear()
1726    }
1727
1728    /// Return prepared derivative cache retention limits.
1729    ///
1730    /// # Examples
1731    ///
1732    /// ```rust
1733    /// use tenferro_ad::EagerRuntime;
1734    /// use tenferro_cpu::CpuBackend;
1735    ///
1736    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1737    /// assert!(ctx.prepared_derivative_cache_limits()?.max_entries().get() > 0);
1738    /// # Ok::<(), tenferro_ad::Error>(())
1739    /// ```
1740    ///
1741    /// # Errors
1742    ///
1743    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
1744    /// derivative cache lock is poisoned.
1745    pub fn prepared_derivative_cache_limits(&self) -> Result<AdTransformCacheLimits> {
1746        Ok(self.lock_prepared_derivative_cache()?.limits())
1747    }
1748
1749    /// Replace prepared derivative cache retention limits.
1750    ///
1751    /// # Examples
1752    ///
1753    /// ```rust
1754    /// use std::num::NonZeroUsize;
1755    /// use tenferro_ad::{AdTransformCacheLimits, EagerRuntime};
1756    /// use tenferro_cpu::CpuBackend;
1757    ///
1758    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1759    /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
1760    /// ctx.set_prepared_derivative_cache_limits(limits)?;
1761    /// assert_eq!(ctx.prepared_derivative_cache_limits()?, limits);
1762    /// # Ok::<(), tenferro_ad::Error>(())
1763    /// ```
1764    ///
1765    /// # Errors
1766    ///
1767    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
1768    /// derivative cache lock is poisoned.
1769    pub fn set_prepared_derivative_cache_limits(
1770        &self,
1771        limits: AdTransformCacheLimits,
1772    ) -> Result<()> {
1773        self.lock_prepared_derivative_cache()?.set_limits(limits);
1774        Ok(())
1775    }
1776
1777    /// Return the extension cache retention limits.
1778    ///
1779    /// # Errors
1780    ///
1781    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the extension
1782    /// cache lock is poisoned.
1783    pub fn extension_cache_limits(&self) -> Result<ExtensionCacheLimits> {
1784        Ok(self.lock_extension_caches()?.limits())
1785    }
1786
1787    /// Replace extension cache retention limits.
1788    ///
1789    /// # Errors
1790    ///
1791    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the extension
1792    /// cache lock is poisoned.
1793    pub fn set_extension_cache_limits(&self, limits: ExtensionCacheLimits) -> Result<()> {
1794        self.lock_extension_caches()?.set_limits(limits);
1795        Ok(())
1796    }
1797
1798    /// Enter one backend execution session and run provider-neutral operations.
1799    ///
1800    /// The callback receives only a lifetime-bound, non-owning backend session.
1801    /// The backend and its engine registration are fixed when the eager runtime
1802    /// is constructed. Extension modules are installed separately and remain
1803    /// available to later extension operations.
1804    ///
1805    /// # Examples
1806    ///
1807    /// ```
1808    /// use tenferro_ad::EagerRuntime;
1809    /// use tenferro_cpu::CpuBackend;
1810    /// use tenferro_tensor::{Tensor, TensorElementwise};
1811    ///
1812    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1813    /// let lhs = Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?;
1814    /// let rhs = Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?;
1815    /// let output = ctx.with_execution_session(|session| {
1816    ///     TensorElementwise::add(session, &lhs, &rhs)
1817    /// })??;
1818    /// assert_eq!(output.as_slice::<f64>()?, &[3.0]);
1819    /// # Ok::<(), tenferro_ad::Error>(())
1820    /// ```
1821    ///
1822    /// # Errors
1823    ///
1824    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the eager backend
1825    /// lock is poisoned. Backend operations retain their typed tensor/backend
1826    /// errors inside the callback result.
1827    pub fn with_execution_session<R: Send>(
1828        &self,
1829        f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
1830    ) -> Result<R> {
1831        let mut backend = self.lock_backend()?;
1832        Ok(backend.with_backend_session(f))
1833    }
1834
1835    // Lock ordering: the eager backend owner is locked first; the
1836    // extension-cache lock is acquired only after it and remains held through
1837    // the borrowed session callback.
1838    /// Run an extension-owned eager operation with a borrowed backend session
1839    /// and the eager runtime's extension cache store.
1840    ///
1841    /// The eager backend owner is locked before the extension-cache lock is
1842    /// acquired. The callback receives an
1843    /// [`tenferro_runtime::ExtensionExecutionContext`] so cache access and
1844    /// backend execution share one lifetime-bound context without exposing the
1845    /// owning eager backend. The backend and its engine registration remain
1846    /// fixed for the eager runtime's lifetime.
1847    ///
1848    /// # Examples
1849    ///
1850    /// ```
1851    /// use tenferro_ad::EagerRuntime;
1852    /// use tenferro_cpu::CpuBackend;
1853    /// use tenferro_tensor::{Tensor, TensorElementwise};
1854    ///
1855    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1856    /// let lhs = Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?;
1857    /// let rhs = Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?;
1858    /// let output = ctx.with_extension_execution_context(|extension_ctx| {
1859    ///     TensorElementwise::add(extension_ctx.backend_mut(), &lhs, &rhs)
1860    /// })??;
1861    /// assert_eq!(output.as_slice::<f64>()?, &[3.0]);
1862    /// # Ok::<(), tenferro_ad::Error>(())
1863    /// ```
1864    ///
1865    /// # Errors
1866    ///
1867    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the eager backend
1868    /// or extension-cache lock is poisoned. Errors returned by the callback
1869    /// remain in its result value.
1870    pub fn with_extension_execution_context<R: Send>(
1871        &self,
1872        f: impl FnOnce(
1873                &mut tenferro_runtime::ExtensionExecutionContext<'_, dyn BackendSession + '_>,
1874            ) -> R
1875            + Send,
1876    ) -> Result<R> {
1877        let mut backend = self.lock_backend()?;
1878        let mut extension_cache_guard = self.lock_extension_caches()?;
1879        let extension_caches: &mut ExtensionCacheStore = &mut extension_cache_guard;
1880        Ok(backend.with_backend_session(move |session| {
1881            let mut extension_ctx =
1882                tenferro_runtime::ExtensionExecutionContext::new(session, extension_caches);
1883            f(&mut extension_ctx)
1884        }))
1885    }
1886
1887    /// Run a prepared extension executor through the runtime-owned erased
1888    /// backend context (the native-context path).
1889    ///
1890    /// This is the sibling of [`Self::with_extension_execution_context`] for
1891    /// prepared operations whose executor does not support the scheduler-owned
1892    /// session but implements the mandatory `execute` bridge. The concrete
1893    /// backend is exposed as an erased context whose type identity matches the
1894    /// executor's binding.
1895    pub(crate) fn with_extension_erased_context<R: Send>(
1896        &self,
1897        f: impl FnOnce(&mut tenferro_runtime::ErasedExecutionContext<'_>, &mut ExtensionCacheStore) -> R
1898            + Send,
1899    ) -> Result<R> {
1900        let mut backend = self.lock_backend()?;
1901        let mut extension_cache_guard = self.lock_extension_caches()?;
1902        let extension_caches: &mut ExtensionCacheStore = &mut extension_cache_guard;
1903        let mut erased = backend.erased_context();
1904        Ok(f(&mut erased, extension_caches))
1905    }
1906
1907    /// Block the current thread until backend work submitted by this eager runtime completes.
1908    ///
1909    /// CPU runtimes return immediately. CUDA and WebGPU runtimes synchronize
1910    /// their current backend work queue.
1911    ///
1912    /// # Examples
1913    ///
1914    /// ```
1915    /// use tenferro_cpu::CpuBackend;
1916    /// use tenferro_ad::EagerRuntime;
1917    ///
1918    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1919    /// ctx.synchronize().unwrap();
1920    /// # Ok::<(), tenferro_ad::Error>(())
1921    /// ```
1922    ///
1923    /// # Errors
1924    ///
1925    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the backend lock is
1926    /// poisoned, or a typed tensor backend error if synchronization fails.
1927    pub fn synchronize(&self) -> Result<()> {
1928        self.lock_backend()?.synchronize().map_err(Error::from)
1929    }
1930
1931    fn exec_outputs_with_runtime<R>(
1932        &self,
1933        lock_backend_section: &'static str,
1934        exec_section: &'static str,
1935        op: &StdTensorOp,
1936        execute: impl FnOnce(&mut EagerBackend, Option<&Runtime>) -> Result<R>,
1937    ) -> Result<R> {
1938        // Lock ordering: eager execution holds the backend lock while standard
1939        // ops run without runtime extension access; extension ops receive the
1940        // runtime so extension cache locks are acquired only from that path.
1941        let mut backend = profile_eager_op_section(lock_backend_section, || self.lock_backend())?;
1942        let runtime = matches!(op, StdTensorOp::Extension(_)).then_some(&self.runtime);
1943        profile_eager_op_section(exec_section, || execute(&mut backend, runtime))
1944    }
1945
1946    pub(crate) fn exec_outputs(&self, op: &StdTensorOp, inputs: &[&Tensor]) -> Result<Vec<Tensor>> {
1947        self.exec_outputs_with_runtime(
1948            "exec_outputs.lock_backend",
1949            "exec_outputs.exec_op",
1950            op,
1951            |backend, runtime| exec_op_on_tensors_with_runtime(op, inputs, backend, runtime),
1952        )
1953    }
1954
1955    pub(crate) fn exec_outputs_read(
1956        &self,
1957        op: &StdTensorOp,
1958        inputs: &[TensorRead<'_>],
1959    ) -> Result<Vec<Tensor>> {
1960        self.exec_outputs_with_runtime(
1961            "exec_outputs_read.lock_backend",
1962            "exec_outputs_read.exec_op",
1963            op,
1964            |backend, runtime| exec_op_on_tensor_reads_with_runtime(op, inputs, backend, runtime),
1965        )
1966    }
1967
1968    #[cfg(test)]
1969    pub(crate) fn exec_standard_graph_outputs(
1970        &self,
1971        graph: &Graph<StdTensorOp>,
1972        initial_data: HashMap<ValueKey<StdTensorOp>, Tensor>,
1973    ) -> Result<EagerGraphExecution> {
1974        let mut backend =
1975            profile_eager_op_section("exec_graph.lock_backend", || self.lock_backend())?;
1976        let mut all_values = initial_data;
1977
1978        profile_eager_op_section("exec_graph.with_backend_session", || {
1979            backend.with_backend_session(|exec| -> Result<()> {
1980                for op_node in graph.operations() {
1981                    let outputs = {
1982                        let input_values = op_node
1983                            .inputs
1984                            .iter()
1985                            .map(|input| {
1986                                let key = match input {
1987                                    ValueRef::Local(local_id) => &graph.values()[*local_id].key,
1988                                    ValueRef::External(key) => key,
1989                                };
1990                                all_values.get(key).ok_or_else(|| {
1991                                    Error::Internal(format!(
1992                                        "standard graph eager execution missing value for {key:?}"
1993                                    ))
1994                                })
1995                            })
1996                            .collect::<Result<Vec<_>>>()?;
1997                        let input_reads = input_values
1998                            .iter()
1999                            .map(|value| TensorRead::from_tensor(value))
2000                            .collect::<Vec<_>>();
2001                        exec_standard_op_on_tensor_reads_in_session(
2002                            &op_node.operation,
2003                            &input_reads,
2004                            exec,
2005                        )?
2006                    };
2007
2008                    if outputs.len() != op_node.outputs.len() {
2009                        return Err(Error::Internal(format!(
2010                            "standard graph eager execution expected {} outputs for {:?}, got {}",
2011                            op_node.outputs.len(),
2012                            op_node.operation,
2013                            outputs.len()
2014                        )));
2015                    }
2016
2017                    for (output_id, output) in op_node.outputs.iter().zip(outputs) {
2018                        let key = graph.values()[*output_id].key.clone();
2019                        all_values.insert(key, output);
2020                    }
2021                }
2022                Ok(())
2023            })
2024        })?;
2025
2026        let outputs = graph
2027            .outputs()
2028            .iter()
2029            .map(|&output_id| {
2030                let key = &graph.values()[output_id].key;
2031                all_values
2032                    .get(key)
2033                    .ok_or_else(|| {
2034                        Error::Internal(format!(
2035                            "standard graph eager execution missing graph output {key:?}"
2036                        ))
2037                    })?
2038                    .duplicate()
2039                    .map_err(Error::from)
2040            })
2041            .collect::<Result<Vec<_>>>()?;
2042
2043        Ok(EagerGraphExecution { outputs })
2044    }
2045
2046    pub(crate) fn try_register_grad_slot(
2047        &self,
2048        key: &ValueKey<StdTensorOp>,
2049        slot: &GradSlot,
2050    ) -> Result<()> {
2051        self.lock_grad_slots()?
2052            .insert(key.clone(), Arc::downgrade(slot));
2053        Ok(())
2054    }
2055
2056    pub(crate) fn try_register_value_record(
2057        &self,
2058        key: &ValueKey<StdTensorOp>,
2059        record: &Arc<EagerTensorRecord>,
2060    ) -> Result<()> {
2061        self.lock_value_records()?
2062            .insert(key.clone(), Arc::downgrade(record));
2063        Ok(())
2064    }
2065
2066    pub(crate) fn value_record(
2067        &self,
2068        key: &ValueKey<StdTensorOp>,
2069    ) -> Result<Option<Arc<EagerTensorRecord>>> {
2070        let mut records = self.lock_value_records()?;
2071        let Some(record) = records.get(key).cloned() else {
2072            return Ok(None);
2073        };
2074        match record.upgrade() {
2075            Some(record) => Ok(Some(record)),
2076            None => {
2077                records.remove(key);
2078                Ok(None)
2079            }
2080        }
2081    }
2082
2083    /// Clear all live gradient slots tracked by this context.
2084    ///
2085    /// This resets the stored gradients to `None` without unregistering the
2086    /// tensors, so future `backward()` calls can accumulate again.
2087    ///
2088    /// # Examples
2089    ///
2090    /// ```
2091    /// use tenferro_cpu::CpuBackend;
2092    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2093    ///
2094    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2095    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx.clone()).unwrap();
2096    /// let y = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![4.0_f64, 5.0, 6.0]).unwrap(), ctx.clone()).unwrap();
2097    /// let loss = x.mul(&y).unwrap().reduce_sum(Some(&[0])).unwrap();
2098    /// let _ = loss.backward().unwrap();
2099    ///
2100    /// ctx.clear_grads()?;
2101    ///
2102    /// assert!(x.grad()?.is_none());
2103    /// assert!(y.grad()?.is_none());
2104    /// # Ok::<(), tenferro_ad::Error>(())
2105    /// ```
2106    ///
2107    /// # Errors
2108    ///
2109    /// Returns [`tenferro_runtime::Error::RuntimeState`] if a gradient-slot
2110    /// lock is poisoned while clearing live gradients.
2111    pub fn clear_grads(&self) -> Result<()> {
2112        let live_slots = {
2113            let mut live_slots = Vec::new();
2114            self.lock_grad_slots()?.retain(|_, slot| {
2115                if let Some(slot) = slot.upgrade() {
2116                    live_slots.push(slot);
2117                    true
2118                } else {
2119                    false
2120                }
2121            });
2122            live_slots
2123        };
2124
2125        let mut poisoned_slot = false;
2126        for slot in live_slots {
2127            match slot.lock() {
2128                Ok(mut current) => {
2129                    *current = None;
2130                }
2131                Err(_) => {
2132                    poisoned_slot = true;
2133                }
2134            }
2135        }
2136        if poisoned_slot {
2137            return Err(Error::runtime_state(
2138                "eager_gradient_slot",
2139                ErrorPhase::Execution,
2140                "lock poisoned",
2141            ));
2142        }
2143        Ok(())
2144    }
2145
2146    /// Import a concrete tensor into this context as an untracked constant.
2147    ///
2148    /// The returned tensor does not participate in gradient tracking.
2149    /// Use this for fixed masks, quadrature weights, physical constants,
2150    /// and other data that should not receive gradients.
2151    ///
2152    /// # Examples
2153    ///
2154    /// ```
2155    /// use tenferro_cpu::CpuBackend;
2156    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2157    ///
2158    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2159    /// let c = ctx.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
2160    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx)?;
2161    /// let z = x.add(&c).unwrap();
2162    ///
2163    /// assert_eq!(z.value()?.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
2164    /// # Ok::<(), tenferro_ad::Error>(())
2165    /// ```
2166    ///
2167    /// # Errors
2168    ///
2169    /// Returns [`tenferro_runtime::Error::RuntimeState`] when metadata cannot
2170    /// be registered or the backend lock is poisoned.
2171    pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
2172        EagerTensor::new_leaf(Arc::clone(self), tensor, false)
2173    }
2174
2175    /// Import a concrete tensor into this context as a trainable variable.
2176    ///
2177    /// The returned tensor participates in gradient tracking; its gradient
2178    /// slot is registered in this context.
2179    ///
2180    /// # Examples
2181    ///
2182    /// ```
2183    /// use tenferro_cpu::CpuBackend;
2184    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2185    ///
2186    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2187    /// let p = ctx.variable_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
2188    /// let loss = p.exp().unwrap().reduce_sum(Some(&[0])).unwrap();
2189    /// let _ = loss.backward().unwrap();
2190    ///
2191    /// let grad = p.grad().unwrap().unwrap();
2192    /// assert_eq!(grad.shape(), &[2]);
2193    /// # Ok::<(), tenferro_ad::Error>(())
2194    /// ```
2195    ///
2196    /// # Errors
2197    ///
2198    /// Returns [`tenferro_runtime::Error::RuntimeState`] when gradient metadata
2199    /// or the eager backend state cannot be registered.
2200    pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
2201        EagerTensor::new_leaf(Arc::clone(self), tensor, true)
2202    }
2203
2204    /// Gradient of a scalar eager output with respect to an eager tensor.
2205    ///
2206    /// Functional eager gradients return ordinary eager tensors and do not
2207    /// write into `grad()` slots. The returned tensor keeps a trace when the
2208    /// derivative computation depends on tracked eager values.
2209    ///
2210    /// # Examples
2211    ///
2212    /// ```
2213    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2214    /// use tenferro_cpu::CpuBackend;
2215    ///
2216    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2217    /// let x = EagerTensor::requires_grad_in(
2218    ///     Tensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap(),
2219    ///     ctx.clone(),
2220    /// )?;
2221    /// let loss = x.mul(&x)?;
2222    /// let dx = ctx.grad(&loss, &x)?;
2223    /// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[6.0]);
2224    /// # Ok::<(), tenferro_ad::Error>(())
2225    /// ```
2226    ///
2227    /// # Errors
2228    ///
2229    /// Returns [`tenferro_runtime::Error::NonScalarGrad`] for a non-scalar
2230    /// output, [`Error::ContextMismatch`] for tensors from another runtime,
2231    /// [`Error::UnsupportedAdRule`] when an AD rule is unavailable, or a typed
2232    /// validation/backend error from eager execution.
2233    pub fn grad(self: &Arc<Self>, output: &EagerTensor, wrt: &EagerTensor) -> Result<EagerTensor> {
2234        self.grad_optional(output, wrt)?
2235            .ok_or_else(|| Error::Internal(format!("grad output is inactive for {:?}", wrt.key)))
2236    }
2237
2238    /// Gradient that returns `None` when `wrt` is inactive.
2239    ///
2240    /// # Examples
2241    ///
2242    /// ```
2243    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2244    /// use tenferro_cpu::CpuBackend;
2245    ///
2246    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2247    /// let x = EagerTensor::requires_grad_in(
2248    ///     Tensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap(),
2249    ///     ctx.clone(),
2250    /// )?;
2251    /// let y = EagerTensor::requires_grad_in(
2252    ///     Tensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap(),
2253    ///     ctx.clone(),
2254    /// )?;
2255    /// let loss = y.mul(&y)?;
2256    /// assert!(ctx.grad_optional(&loss, &x)?.is_none());
2257    /// # Ok::<(), tenferro_ad::Error>(())
2258    /// ```
2259    ///
2260    /// # Errors
2261    ///
2262    /// Returns [`tenferro_runtime::Error::NonScalarGrad`] for a non-scalar
2263    /// output, [`Error::ContextMismatch`] for a foreign runtime, or a typed
2264    /// validation/backend/runtime-state error from eager execution.
2265    pub fn grad_optional(
2266        self: &Arc<Self>,
2267        output: &EagerTensor,
2268        wrt: &EagerTensor,
2269    ) -> Result<Option<EagerTensor>> {
2270        if !output.shape().is_empty() {
2271            return Err(Error::NonScalarGrad {
2272                shape: output.shape().to_vec(),
2273            });
2274        }
2275
2276        let value = output.to_tensor()?;
2277        let seed = {
2278            let mut backend = self.lock_backend()?;
2279            one_like_tensor(&value, &mut *backend)?
2280        };
2281        let seed = EagerTensor::new_result(Arc::clone(self), eager_val_key(), seed, false, None)?;
2282        self.vjp_optional(output, wrt, &seed)
2283    }
2284
2285    /// Reverse-mode vector-Jacobian product for eager tensors.
2286    ///
2287    /// # Examples
2288    ///
2289    /// ```
2290    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2291    /// use tenferro_cpu::CpuBackend;
2292    ///
2293    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2294    /// let x = EagerTensor::requires_grad_in(
2295    ///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
2296    ///     ctx.clone(),
2297    /// )?;
2298    /// let y = x.mul(&x)?;
2299    /// let seed = EagerTensor::from_tensor_in(
2300    ///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 1.0]).unwrap(),
2301    ///     ctx.clone(),
2302    /// )?;
2303    /// let dx = ctx.vjp(&y, &x, &seed)?;
2304    /// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
2305    /// # Ok::<(), tenferro_ad::Error>(())
2306    /// ```
2307    ///
2308    /// # Errors
2309    ///
2310    /// Returns [`Error::ContextMismatch`] for tensors from different eager
2311    /// runtimes, [`Error::Validation`] when the cotangent shape or dtype does
2312    /// not match the output, [`Error::UnsupportedAdRule`] when a rule is not
2313    /// registered, or a typed backend/runtime-state error.
2314    pub fn vjp(
2315        self: &Arc<Self>,
2316        output: &EagerTensor,
2317        wrt: &EagerTensor,
2318        cotangent: &EagerTensor,
2319    ) -> Result<EagerTensor> {
2320        self.vjp_optional(output, wrt, cotangent)?
2321            .ok_or_else(|| Error::Internal(format!("vjp output is inactive for {:?}", wrt.key)))
2322    }
2323
2324    /// Reverse-mode vector-Jacobian product that returns `None` for inactive inputs.
2325    ///
2326    /// # Examples
2327    ///
2328    /// ```
2329    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2330    /// use tenferro_cpu::CpuBackend;
2331    ///
2332    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2333    /// let x = EagerTensor::requires_grad_in(
2334    ///     Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
2335    ///     ctx.clone(),
2336    /// )?;
2337    /// let y = EagerTensor::requires_grad_in(
2338    ///     Tensor::from_vec_col_major(vec![1], vec![4.0_f64]).unwrap(),
2339    ///     ctx.clone(),
2340    /// )?;
2341    /// let seed = EagerTensor::from_tensor_in(
2342    ///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
2343    ///     ctx.clone(),
2344    /// )?;
2345    /// let loss = y.mul(&y)?;
2346    /// assert!(ctx.vjp_optional(&loss, &x, &seed)?.is_none());
2347    /// # Ok::<(), tenferro_ad::Error>(())
2348    /// ```
2349    ///
2350    /// # Errors
2351    ///
2352    /// Returns [`Error::ContextMismatch`] for tensors from different eager
2353    /// runtimes, [`Error::Validation`] when the cotangent shape or dtype does
2354    /// not match the output, [`Error::UnsupportedAdRule`] when a rule is not
2355    /// registered, or a typed backend/runtime-state error.
2356    pub fn vjp_optional(
2357        self: &Arc<Self>,
2358        output: &EagerTensor,
2359        wrt: &EagerTensor,
2360        cotangent: &EagerTensor,
2361    ) -> Result<Option<EagerTensor>> {
2362        validate_same_runtime(self, output, "vjp output")?;
2363        validate_same_runtime(self, wrt, "vjp wrt")?;
2364        validate_same_runtime(self, cotangent, "vjp cotangent")?;
2365        validate_seed_tensor("vjp", output, cotangent)?;
2366        // Unification 7: semantic path is the only VJP path.
2367        match semantic_eager_vjp_optional(self, output, wrt, cotangent)? {
2368            Some(result) => Ok(result),
2369            None => Ok(None),
2370        }
2371    }
2372
2373    /// Forward-mode Jacobian-vector product for eager tensors.
2374    ///
2375    /// # Examples
2376    ///
2377    /// ```
2378    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2379    /// use tenferro_cpu::CpuBackend;
2380    ///
2381    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2382    /// let x = EagerTensor::requires_grad_in(
2383    ///     Tensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap(),
2384    ///     ctx.clone(),
2385    /// )?;
2386    /// let tangent = EagerTensor::from_tensor_in(
2387    ///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
2388    ///     ctx.clone(),
2389    /// )?;
2390    /// let y = x.mul(&x)?;
2391    /// let dy = ctx.jvp(&y, &x, &tangent)?;
2392    /// assert_eq!(dy.value()?.as_slice::<f64>().unwrap(), &[6.0]);
2393    /// # Ok::<(), tenferro_ad::Error>(())
2394    /// ```
2395    ///
2396    /// # Errors
2397    ///
2398    /// Returns [`Error::ContextMismatch`] for tensors from different eager
2399    /// runtimes, [`Error::Validation`] when the tangent shape or dtype does not
2400    /// match `wrt`, [`Error::UnsupportedAdRule`] when a rule is unavailable, or
2401    /// a typed backend/runtime-state error.
2402    pub fn jvp(
2403        self: &Arc<Self>,
2404        output: &EagerTensor,
2405        wrt: &EagerTensor,
2406        tangent: &EagerTensor,
2407    ) -> Result<EagerTensor> {
2408        self.jvp_optional(output, wrt, tangent)?
2409            .ok_or_else(|| Error::Internal(format!("jvp output is inactive for {:?}", wrt.key)))
2410    }
2411
2412    /// Forward-mode Jacobian-vector product that returns `None` for inactive outputs.
2413    ///
2414    /// # Examples
2415    ///
2416    /// ```
2417    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2418    /// use tenferro_cpu::CpuBackend;
2419    ///
2420    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
2421    /// let x = EagerTensor::requires_grad_in(
2422    ///     Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
2423    ///     ctx.clone(),
2424    /// )?;
2425    /// let y = EagerTensor::requires_grad_in(
2426    ///     Tensor::from_vec_col_major(vec![1], vec![4.0_f64]).unwrap(),
2427    ///     ctx.clone(),
2428    /// )?;
2429    /// let tangent = EagerTensor::from_tensor_in(
2430    ///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
2431    ///     ctx.clone(),
2432    /// )?;
2433    /// let loss = y.mul(&y)?;
2434    /// assert!(ctx.jvp_optional(&loss, &x, &tangent)?.is_none());
2435    /// # Ok::<(), tenferro_ad::Error>(())
2436    /// ```
2437    ///
2438    /// # Errors
2439    ///
2440    /// Returns [`Error::ContextMismatch`] for tensors from different eager
2441    /// runtimes, [`Error::Validation`] when the tangent shape or dtype does not
2442    /// match `wrt`, [`Error::UnsupportedAdRule`] when a rule is unavailable, or
2443    /// a typed backend/runtime-state error.
2444    pub fn jvp_optional(
2445        self: &Arc<Self>,
2446        output: &EagerTensor,
2447        wrt: &EagerTensor,
2448        tangent: &EagerTensor,
2449    ) -> Result<Option<EagerTensor>> {
2450        validate_same_runtime(self, output, "jvp output")?;
2451        validate_same_runtime(self, wrt, "jvp wrt")?;
2452        validate_same_runtime(self, tangent, "jvp tangent")?;
2453        validate_seed_tensor("jvp", wrt, tangent)?;
2454        // Unification 7: semantic path is the only JVP path.
2455        match semantic_eager_jvp_optional(self, output, wrt, tangent)? {
2456            Some(result) => Ok(result),
2457            None => Ok(None),
2458        }
2459    }
2460
2461    fn store_grads(
2462        &self,
2463        cotangents: &HashMap<ValueKey<StdTensorOp>, Tensor>,
2464        backend: &mut EagerBackend,
2465    ) -> Result<()> {
2466        let mut updates = Vec::new();
2467
2468        {
2469            let mut slots = self.lock_grad_slots()?;
2470            slots.retain(|key, slot| {
2471                let Some(slot) = slot.upgrade() else {
2472                    return false;
2473                };
2474
2475                if let Some(incoming) = cotangents.get(key) {
2476                    updates.push((slot, incoming));
2477                }
2478
2479                true
2480            });
2481        }
2482
2483        for (slot, incoming) in updates {
2484            let mut current = slot.lock().map_err(|_| {
2485                Error::runtime_state(
2486                    "eager_gradient_slot",
2487                    ErrorPhase::Execution,
2488                    "lock poisoned",
2489                )
2490            })?;
2491            let next = match current.as_ref() {
2492                Some(existing) => {
2493                    let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
2494                    let incoming_read = TensorRead::from_tensor(incoming);
2495                    let tensor = backend
2496                        .with_backend_session(|session| {
2497                            session.add_read(existing_read, incoming_read)
2498                        })
2499                        .map_err(Error::from)?;
2500                    AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
2501                }
2502                None => {
2503                    let duplicate = backend
2504                        .with_backend_session(|session| {
2505                            session.to_contiguous_read(TensorRead::from_tensor(incoming))
2506                        })
2507                        .map_err(Error::from)?;
2508                    AdValueRecord::from_tensor(duplicate, "EagerRuntime::store_grads")?
2509                }
2510            };
2511            *current = Some(next);
2512        }
2513
2514        Ok(())
2515    }
2516}
2517
2518#[derive(Clone, Debug, PartialEq, Eq, Hash)]
2519struct PreparedDerivativeCacheKey {
2520    semantic_fingerprint: SemanticFingerprint,
2521    runtime_epoch: RuntimeEpoch,
2522    wrt_input_index: usize,
2523    input_metadata: Box<[ProgramValueMetadata]>,
2524}
2525
2526/// Cached prepared derivative: program + index metadata.
2527#[derive(Debug)]
2528struct PreparedDerivative {
2529    program: Arc<CompiledGraph>,
2530    prepared: Arc<PreparedCompiledGraph>,
2531    seed_input_index: usize,
2532    derivative_output_index: usize,
2533}
2534
2535#[derive(Debug)]
2536struct PreparedDerivativeCache {
2537    limits: AdTransformCacheLimits,
2538    entries: LruCache<PreparedDerivativeCacheKey, PreparedDerivativeCacheEntry>,
2539    stats: CacheStats,
2540}
2541
2542impl PreparedDerivativeCache {
2543    fn limits(&self) -> AdTransformCacheLimits {
2544        self.limits
2545    }
2546
2547    fn set_limits(&mut self, limits: AdTransformCacheLimits) {
2548        self.limits = limits;
2549        self.evict_to_limits();
2550    }
2551
2552    fn clear(&mut self) {
2553        let clears = self.stats.clears.saturating_add(1);
2554        self.entries.clear();
2555        self.stats = CacheStats {
2556            clears,
2557            ..CacheStats::empty()
2558        };
2559    }
2560
2561    fn stats(&self) -> CacheStats {
2562        self.stats
2563    }
2564
2565    fn get(&mut self, key: &PreparedDerivativeCacheKey) -> Option<Arc<PreparedDerivative>> {
2566        match self.entries.get(key) {
2567            Some(entry) => {
2568                self.stats.hits = self.stats.hits.saturating_add(1);
2569                Some(Arc::clone(&entry.value))
2570            }
2571            None => {
2572                self.stats.misses = self.stats.misses.saturating_add(1);
2573                None
2574            }
2575        }
2576    }
2577
2578    fn insert(&mut self, key: PreparedDerivativeCacheKey, value: Arc<PreparedDerivative>) {
2579        let retained_bytes = prepared_derivative_cache_entry_retained_bytes(&key, value.as_ref());
2580        let entry = PreparedDerivativeCacheEntry {
2581            value,
2582            retained_bytes,
2583        };
2584        self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
2585        if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
2586            self.stats.retained_bytes = self
2587                .stats
2588                .retained_bytes
2589                .saturating_sub(old_entry.retained_bytes);
2590        }
2591        self.stats.entries = self.entries.len();
2592        self.evict_to_limits();
2593    }
2594
2595    fn evict_to_limits(&mut self) {
2596        while self.entries.len() > self.limits.max_entries().get()
2597            || self
2598                .limits
2599                .max_retained_bytes()
2600                .is_some_and(|limit| self.stats.retained_bytes > limit.get())
2601        {
2602            let Some((_key, entry)) = self.entries.pop_lru() else {
2603                break;
2604            };
2605            self.stats.retained_bytes = self
2606                .stats
2607                .retained_bytes
2608                .saturating_sub(entry.retained_bytes);
2609            self.stats.evictions = self.stats.evictions.saturating_add(1);
2610        }
2611        self.stats.entries = self.entries.len();
2612    }
2613}
2614
2615impl Default for PreparedDerivativeCache {
2616    fn default() -> Self {
2617        Self {
2618            limits: AdTransformCacheLimits::default(),
2619            entries: LruCache::unbounded(),
2620            stats: CacheStats::empty(),
2621        }
2622    }
2623}
2624
2625#[derive(Debug)]
2626struct PreparedDerivativeCacheEntry {
2627    value: Arc<PreparedDerivative>,
2628    retained_bytes: usize,
2629}
2630
2631fn prepared_derivative_cache_entry_retained_bytes(
2632    key: &PreparedDerivativeCacheKey,
2633    value: &PreparedDerivative,
2634) -> usize {
2635    size_of::<PreparedDerivativeCacheKey>()
2636        .saturating_add(
2637            key.input_metadata
2638                .len()
2639                .saturating_mul(size_of::<ProgramValueMetadata>()),
2640        )
2641        .saturating_add(size_of::<PreparedDerivative>())
2642        .saturating_add(compiled_graph_retained_bytes(value.program.as_ref()))
2643        .saturating_add(prepared_compiled_graph_retained_bytes(
2644            value.prepared.as_ref(),
2645            value.program.as_ref(),
2646        ))
2647}
2648
2649fn prepared_compiled_graph_retained_bytes(
2650    prepared: &PreparedCompiledGraph,
2651    derivative_program: &CompiledGraph,
2652) -> usize {
2653    size_of_val(prepared).saturating_add(compiled_graph_retained_bytes(derivative_program))
2654}
2655
2656fn compiled_graph_retained_bytes(program: &CompiledGraph) -> usize {
2657    size_of::<CompiledGraph>()
2658        .saturating_add(size_of_val(program.input_keys()))
2659        .saturating_add(program.bindings().len().saturating_mul(size_of::<usize>()))
2660        .saturating_add(semantic_program_retained_bytes(program.program()))
2661}
2662
2663fn semantic_program_retained_bytes(program: &SemanticProgram) -> usize {
2664    size_of::<SemanticProgram>()
2665        .saturating_add(size_of_val(program.inputs()))
2666        .saturating_add(size_of_val(program.outputs()))
2667        .saturating_add(
2668            program
2669                .operations()
2670                .len()
2671                .saturating_mul(size_of::<usize>()),
2672        )
2673        .saturating_add(
2674            program
2675                .shape_guards()
2676                .len()
2677                .saturating_mul(size_of::<usize>()),
2678        )
2679}
2680
2681fn semantic_eager_vjp_optional(
2682    ctx: &Arc<EagerRuntime>,
2683    output: &EagerTensor,
2684    wrt: &EagerTensor,
2685    cotangent: &EagerTensor,
2686) -> Result<Option<Option<EagerTensor>>> {
2687    if !eager_semantic_vjp_enabled() {
2688        return Ok(None);
2689    }
2690    let (Some(raw_output_trace), Some(wrt_trace)) =
2691        (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
2692    else {
2693        return Ok(None);
2694    };
2695    let Some(wrt_key) = wrt_trace.input_key() else {
2696        return Ok(None);
2697    };
2698    if !raw_output_trace.has_attached_input_key(&wrt_key) {
2699        return Ok(None);
2700    }
2701
2702    // First AD request on this output: run the deferred graph analysis over
2703    // the whole raw carrier chain once (metadata registration + constraint
2704    // scopes), so `compile_ad_source` sees the same analyzed graph the eager
2705    // forward used to append.
2706    let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
2707
2708    // First compile the trace to get bindings and wrt_input_index.
2709    // (The compile step is needed even for cache hits to extract tensor bindings.)
2710    let mut compiler = GraphCompiler::new();
2711    let source = compile_ad_source(&mut compiler, &output_trace)?;
2712    if source.output_count() != 1
2713        || source.input_keys().len() != source.input_count()
2714        || source.bindings().len() != source.input_count()
2715    {
2716        return Ok(None);
2717    }
2718    let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
2719        return Ok(None);
2720    };
2721
2722    // S2: check prepared-derivative cache before AD transform + compile_frozen.
2723    let cache_key = PreparedDerivativeCacheKey {
2724        semantic_fingerprint: source.program().semantic_fingerprint(),
2725        runtime_epoch: ctx.runtime.epoch().map_err(|source| {
2726            Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::Execution, source)
2727        })?,
2728        wrt_input_index,
2729        input_metadata: source.frozen_program().input_metadata_with_bound_shapes(),
2730    };
2731    let prepared = { ctx.lock_prepared_derivative_cache()?.get(&cache_key) };
2732    let (seed_input_index, derivative_output_index, derivative_program, prepared_runtime) =
2733        if let Some(prepared) = prepared {
2734            (
2735                prepared.seed_input_index,
2736                prepared.derivative_output_index,
2737                Arc::clone(&prepared.program),
2738                Some(Arc::clone(&prepared.prepared)),
2739            )
2740        } else {
2741            let mut active_inputs = vec![false; source.input_count()];
2742            if let Some(active) = active_inputs.get_mut(wrt_input_index) {
2743                *active = true;
2744            } else {
2745                return Ok(None);
2746            }
2747            let active_outputs = vec![true; source.output_count()];
2748            let ad = AdContext::with_rules_and_transform_cache(
2749                ctx.semantic_extension_rules.clone(),
2750                Arc::clone(&ctx.ad_transform_cache),
2751            );
2752            let derivative = ad
2753                .vjp_program(source.frozen_program(), &active_inputs, &active_outputs)
2754                .map_err(|source| {
2755                    Error::runtime_state_source(
2756                        "semantic_eager_vjp",
2757                        ErrorPhase::GraphBuild,
2758                        source,
2759                    )
2760                })?;
2761            let seed_input_index = derivative
2762                .derivative_input_indices()
2763                .first()
2764                .copied()
2765                .flatten();
2766            let derivative_output_index = derivative
2767                .derivative_output_indices()
2768                .get(wrt_input_index)
2769                .copied()
2770                .flatten();
2771            let (Some(seed_input_index), Some(derivative_output_index)) =
2772                (seed_input_index, derivative_output_index)
2773            else {
2774                return Ok(Some(None));
2775            };
2776            let program = Arc::new(compiler.compile_frozen_program(derivative.frozen())?);
2777            (seed_input_index, derivative_output_index, program, None)
2778        };
2779
2780    let cotangent_tensor = Arc::new(RetainedValue::from_tensor(cotangent.to_tensor()?));
2781    let input_count = derivative_program.input_count();
2782    let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
2783    for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
2784        let Some(slot) = owned_inputs.get_mut(source_input_index) else {
2785            return Err(Error::Internal(format!(
2786                "semantic eager VJP derivative program has no primal input slot {source_input_index}"
2787            )));
2788        };
2789        *slot = Some(copy_value_for_runtime(ctx, tensor)?);
2790    }
2791    let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
2792        return Err(Error::Internal(format!(
2793            "semantic eager VJP seed input index {seed_input_index} is outside {} inputs",
2794            owned_inputs.len()
2795        )));
2796    };
2797    *slot = Some(copy_value_for_runtime(ctx, cotangent_tensor.as_ref())?);
2798    let input_refs = owned_inputs
2799        .iter()
2800        .enumerate()
2801        .map(|(index, tensor)| {
2802            tensor.as_ref().ok_or_else(|| {
2803                Error::Internal(format!(
2804                    "semantic eager VJP derivative input {index} was not populated"
2805                ))
2806            })
2807        })
2808        .collect::<Result<Vec<_>>>()?;
2809    let prepared_runtime = if let Some(prepared_runtime) = prepared_runtime {
2810        prepared_runtime
2811    } else {
2812        let prepared_runtime = Arc::new(
2813            ctx.runtime
2814                .prepare_compiled(&derivative_program, &input_refs)?,
2815        );
2816        let entry = Arc::new(PreparedDerivative {
2817            program: Arc::clone(&derivative_program),
2818            prepared: Arc::clone(&prepared_runtime),
2819            seed_input_index,
2820            derivative_output_index,
2821        });
2822        ctx.lock_prepared_derivative_cache()?
2823            .insert(cache_key, entry);
2824        prepared_runtime
2825    };
2826    let outputs = ctx.runtime.run_prepared(&prepared_runtime, &input_refs)?;
2827    let output_count = outputs.len();
2828    let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
2829        return Err(Error::Internal(format!(
2830            "semantic eager VJP derivative output index {derivative_output_index} is outside {} outputs",
2831            output_count
2832        )));
2833    };
2834    let cotangent_trace =
2835        TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&cotangent_tensor))?;
2836    let semantic_trace = derivative_trace_from_frozen_program(
2837        &source,
2838        derivative_program.frozen_program(),
2839        derivative_output_index,
2840        &[(seed_input_index, Arc::clone(&cotangent_tensor))],
2841        &[&output_trace, wrt_trace, &cotangent_trace],
2842        None,
2843        "semantic_eager_vjp",
2844    )?;
2845
2846    #[cfg(test)]
2847    EAGER_SEMANTIC_VJP_EXECUTIONS.fetch_add(1, Ordering::Relaxed);
2848
2849    Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
2850        Arc::clone(ctx),
2851        eager_val_key(),
2852        result,
2853        true,
2854        None,
2855        Some(semantic_trace),
2856    )?)))
2857}
2858
2859fn semantic_eager_jvp_optional(
2860    ctx: &Arc<EagerRuntime>,
2861    output: &EagerTensor,
2862    wrt: &EagerTensor,
2863    tangent: &EagerTensor,
2864) -> Result<Option<Option<EagerTensor>>> {
2865    if !eager_semantic_vjp_enabled() {
2866        return Ok(None);
2867    }
2868    let (Some(raw_output_trace), Some(wrt_trace)) =
2869        (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
2870    else {
2871        return Ok(None);
2872    };
2873    let Some(wrt_key) = wrt_trace.input_key() else {
2874        return Ok(None);
2875    };
2876    if !raw_output_trace.has_attached_input_key(&wrt_key) {
2877        return Ok(None);
2878    }
2879
2880    // First AD request on this output: run the deferred graph analysis once
2881    // over the whole raw carrier chain before compiling.
2882    let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
2883
2884    let mut compiler = GraphCompiler::new();
2885    let source = compile_ad_source(&mut compiler, &output_trace)?;
2886    if source.output_count() != 1
2887        || source.input_keys().len() != source.input_count()
2888        || source.bindings().len() != source.input_count()
2889    {
2890        return Ok(None);
2891    }
2892    let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
2893        return Ok(None);
2894    };
2895
2896    let mut active_inputs = vec![false; source.input_count()];
2897    if let Some(active) = active_inputs.get_mut(wrt_input_index) {
2898        *active = true;
2899    } else {
2900        return Ok(None);
2901    }
2902    let ad = AdContext::with_rules_and_transform_cache(
2903        ctx.semantic_extension_rules.clone(),
2904        Arc::clone(&ctx.ad_transform_cache),
2905    );
2906    let derivative = ad
2907        .jvp_program(source.frozen_program(), &active_inputs)
2908        .map_err(|source| {
2909            Error::runtime_state_source("semantic_eager_jvp", ErrorPhase::GraphBuild, source)
2910        })?;
2911    // derivative_input_indices maps source input → derivative seed input.
2912    let Some(seed_input_index) = derivative
2913        .derivative_input_indices()
2914        .get(wrt_input_index)
2915        .copied()
2916        .flatten()
2917    else {
2918        return Ok(Some(None));
2919    };
2920    // derivative_output_indices maps source output → derivative output.
2921    // There is always exactly one source output (guarded above).
2922    let Some(derivative_output_index) = derivative
2923        .derivative_output_indices()
2924        .first()
2925        .copied()
2926        .flatten()
2927    else {
2928        return Ok(Some(None));
2929    };
2930
2931    let derivative_program = compiler.compile_frozen_program(derivative.frozen())?;
2932    let tangent_tensor = Arc::new(RetainedValue::from_tensor(tangent.to_tensor()?));
2933    let input_count = derivative_program.input_count();
2934    let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
2935    for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
2936        let Some(slot) = owned_inputs.get_mut(source_input_index) else {
2937            return Err(Error::Internal(format!(
2938                "semantic eager JVP derivative program has no primal input slot {source_input_index}"
2939            )));
2940        };
2941        *slot = Some(copy_value_for_runtime(ctx, tensor)?);
2942    }
2943    let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
2944        return Err(Error::Internal(format!(
2945            "semantic eager JVP seed input index {seed_input_index} is outside {} inputs",
2946            owned_inputs.len()
2947        )));
2948    };
2949    *slot = Some(copy_value_for_runtime(ctx, tangent_tensor.as_ref())?);
2950    let input_refs = owned_inputs
2951        .iter()
2952        .enumerate()
2953        .map(|(index, tensor)| {
2954            tensor.as_ref().ok_or_else(|| {
2955                Error::Internal(format!(
2956                    "semantic eager JVP derivative input {index} was not populated"
2957                ))
2958            })
2959        })
2960        .collect::<Result<Vec<_>>>()?;
2961    let outputs = ctx.runtime.run_compiled(&derivative_program, &input_refs)?;
2962    let output_count = outputs.len();
2963    let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
2964        return Err(Error::Internal(format!(
2965            "semantic eager JVP derivative output index {derivative_output_index} is outside {} outputs",
2966            output_count
2967        )));
2968    };
2969    let tangent_trace =
2970        TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&tangent_tensor))?;
2971    let semantic_trace = derivative_trace_from_frozen_program(
2972        &source,
2973        derivative.frozen(),
2974        derivative_output_index,
2975        &[(seed_input_index, Arc::clone(&tangent_tensor))],
2976        &[&output_trace, wrt_trace, &tangent_trace],
2977        None,
2978        "semantic_eager_jvp",
2979    )?;
2980
2981    Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
2982        Arc::clone(ctx),
2983        eager_val_key(),
2984        result,
2985        true,
2986        None,
2987        Some(semantic_trace),
2988    )?)))
2989}
2990
2991fn validate_same_runtime(
2992    runtime: &Arc<EagerRuntime>,
2993    tensor: &EagerTensor,
2994    role: &'static str,
2995) -> Result<()> {
2996    if tensor.ctx_id() != runtime.id() {
2997        return Err(Error::ContextMismatch {
2998            lhs: runtime.id(),
2999            rhs: tensor.ctx_id(),
3000        });
3001    }
3002    let _ = role;
3003    Ok(())
3004}
3005
3006fn copy_value_for_runtime(ctx: &EagerRuntime, value: &RetainedValue) -> Result<Tensor> {
3007    let read = value.tensor_read().map_err(|error| {
3008        Error::runtime_state_source("copy_value_for_runtime", ErrorPhase::Execution, error)
3009    })?;
3010    ctx.with_execution_session(|session| session.to_contiguous_read(read))?
3011        .map_err(Error::from)
3012}
3013
3014fn validate_seed_tensor(op: &'static str, primal: &EagerTensor, seed: &EagerTensor) -> Result<()> {
3015    if primal.dtype() != seed.dtype() {
3016        return Err(
3017            tenferro_tensor::Error::dtype_mismatch(op, primal.dtype(), seed.dtype()).into(),
3018        );
3019    }
3020    if primal.shape() != seed.shape() {
3021        return Err(
3022            tenferro_tensor::Error::shape_mismatch(op, primal.shape(), seed.shape()).into(),
3023        );
3024    }
3025    Ok(())
3026}
3027
3028/// Eager tensor with reverse-mode autodiff over concrete tensor values.
3029///
3030/// This executes each primitive immediately and records a lightweight reverse
3031/// DAG for `backward()`. Gradients accumulate across repeated `backward()`
3032/// calls until they are cleared explicitly.
3033///
3034/// # Examples
3035///
3036/// ```
3037/// use tenferro_cpu::CpuBackend;
3038/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3039///
3040/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3041/// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx)?;
3042/// let loss = x.mul(&x).unwrap().reduce_sum(Some(&[0])).unwrap();
3043/// let _cotangents = loss.backward().unwrap();
3044/// let loss = x.mul(&x).unwrap().reduce_sum(Some(&[0])).unwrap();
3045/// let _cotangents = loss.backward().unwrap();
3046///
3047/// assert_eq!(x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(), &[4.0, 8.0, 12.0]);
3048/// x.clear_grad();
3049///
3050/// assert!(x.grad().unwrap().is_none());
3051/// # Ok::<(), tenferro_ad::Error>(())
3052/// ```
3053#[derive(Clone)]
3054pub struct EagerTensor {
3055    pub(crate) key: ValueKey<StdTensorOp>,
3056    pub(crate) trace: Option<EagerTrace>,
3057    pub(crate) semantic_trace: Option<TracedTensor>,
3058    pub(crate) requires_grad: bool,
3059    grad_slot: GradSlot,
3060    pub(crate) ctx: Arc<EagerRuntime>,
3061    _record: Arc<EagerTensorRecord>,
3062}
3063
3064pub(crate) struct EagerTensorRecord {
3065    value: Arc<AdValueRecord>,
3066    key: ValueKey<StdTensorOp>,
3067    trace: Option<EagerTrace>,
3068    semantic_trace: Option<TracedTensor>,
3069    requires_grad: bool,
3070    grad_slot: GradSlot,
3071    ctx: Arc<EagerRuntime>,
3072}
3073
3074struct EagerTensorParts {
3075    ctx: Arc<EagerRuntime>,
3076    key: ValueKey<StdTensorOp>,
3077    requires_grad: bool,
3078    trace: Option<EagerTrace>,
3079    semantic_trace: Option<TracedTensor>,
3080    value: Arc<AdValueRecord>,
3081    register_value: bool,
3082}
3083
3084impl fmt::Debug for EagerTensor {
3085    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3086        f.debug_struct("EagerTensor")
3087            .field("dtype", &self.dtype())
3088            .field("shape", &self.shape())
3089            .field("key", &self.key)
3090            .field("requires_grad", &self.requires_grad)
3091            .field("has_trace", &self.trace.is_some())
3092            .field("has_semantic_trace", &self.semantic_trace.is_some())
3093            .field("ctx_id", &self.ctx_id())
3094            .finish_non_exhaustive()
3095    }
3096}
3097
3098impl EagerTensor {
3099    /// Create an untracked eager tensor inside an existing eager context.
3100    ///
3101    /// # Examples
3102    ///
3103    /// ```
3104    /// use tenferro_cpu::CpuBackend;
3105    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3106    ///
3107    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3108    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
3109    ///
3110    /// assert_eq!(x.value()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
3111    /// # Ok::<(), tenferro_ad::Error>(())
3112    /// ```
3113    ///
3114    /// # Errors
3115    ///
3116    /// Returns [`tenferro_runtime::Error::RuntimeState`] when metadata cannot
3117    /// be registered in the target context, or a typed tensor/backend error
3118    /// while materializing the source value.
3119    pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
3120        Self::new_leaf(ctx, tensor, false)
3121    }
3122
3123    /// Create an untracked eager tensor from compact column-major data inside
3124    /// an existing eager runtime.
3125    ///
3126    /// # Errors
3127    ///
3128    /// Returns [`Error::TensorRuntime`] with
3129    /// [`tenferro_tensor::ValidationError::ShapeMismatch`] when the shape and
3130    /// data length disagree, or with
3131    /// [`tenferro_tensor::ValidationError::IntegerOverflow`] when shape
3132    /// arithmetic overflows. Returns [`Error::RuntimeState`] when eager
3133    /// metadata cannot be registered.
3134    pub fn from_vec_col_major_in<T: TensorScalar>(
3135        shape: impl IntoShapeVec,
3136        data: Vec<T>,
3137        ctx: Arc<EagerRuntime>,
3138    ) -> Result<Self> {
3139        Self::from_tensor_in(Tensor::from_vec_col_major(shape, data)?, ctx)
3140    }
3141
3142    /// Create a tracked eager leaf inside an existing eager context.
3143    ///
3144    /// # Examples
3145    ///
3146    /// ```
3147    /// use tenferro_cpu::CpuBackend;
3148    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3149    ///
3150    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3151    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
3152    ///
3153    /// assert!(x.grad().unwrap().is_none());
3154    /// # Ok::<(), tenferro_ad::Error>(())
3155    /// ```
3156    ///
3157    /// # Errors
3158    ///
3159    /// Returns [`tenferro_runtime::Error::RuntimeState`] when gradient metadata
3160    /// cannot be registered in the target context, or a typed tensor/backend
3161    /// error while creating the leaf.
3162    pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
3163        Self::new_leaf(ctx, tensor, true)
3164    }
3165
3166    pub(crate) fn new_leaf(
3167        ctx: Arc<EagerRuntime>,
3168        tensor: Tensor,
3169        requires_grad: bool,
3170    ) -> Result<Self> {
3171        let key = eager_val_key();
3172        let semantic_tensor = ctx
3173            .with_execution_session(|session| {
3174                session.to_contiguous_read(TensorRead::from_tensor(&tensor))
3175            })?
3176            .map_err(Error::from)?;
3177        let semantic_value = Arc::new(RetainedValue::from_tensor(semantic_tensor));
3178        let semantic_trace = TracedTensor::from_shared_tensor_value_symbolic_shape(semantic_value)?;
3179        // Deferred materialization: the per-op/leaf global-metadata registry
3180        // write for the eager tensor key was unreadable after the semantic
3181        // trace became the sole AD carrier, so it is dropped.
3182        // ponytail: leaf input-key metadata is still registered by
3183        // `from_shared_tensor_value_symbolic_shape`; the eager-key entry was
3184        // vestigial and removed. Add back only if something reads it.
3185        let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_leaf")?;
3186        Self::from_parts(EagerTensorParts {
3187            ctx,
3188            key,
3189            requires_grad,
3190            trace: None,
3191            semantic_trace: Some(semantic_trace),
3192            value,
3193            register_value: true,
3194        })
3195    }
3196
3197    pub(crate) fn new_result(
3198        ctx: Arc<EagerRuntime>,
3199        key: ValueKey<StdTensorOp>,
3200        tensor: Tensor,
3201        requires_grad: bool,
3202        trace: Option<EagerTrace>,
3203    ) -> Result<Self> {
3204        Self::new_result_with_semantic_trace(ctx, key, tensor, requires_grad, trace, None)
3205    }
3206
3207    pub(crate) fn new_result_with_semantic_trace(
3208        ctx: Arc<EagerRuntime>,
3209        key: ValueKey<StdTensorOp>,
3210        tensor: Tensor,
3211        requires_grad: bool,
3212        trace: Option<EagerTrace>,
3213        semantic_trace: Option<TracedTensor>,
3214    ) -> Result<Self> {
3215        let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_result")?;
3216        Self::from_parts(EagerTensorParts {
3217            ctx,
3218            key,
3219            requires_grad,
3220            trace,
3221            semantic_trace,
3222            value,
3223            register_value: true,
3224        })
3225    }
3226
3227    pub(crate) fn new_unregistered_result_with_semantic_trace(
3228        ctx: Arc<EagerRuntime>,
3229        key: ValueKey<StdTensorOp>,
3230        tensor: Tensor,
3231        requires_grad: bool,
3232        trace: Option<EagerTrace>,
3233        semantic_trace: Option<TracedTensor>,
3234    ) -> Result<Self> {
3235        let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_unregistered_result")?;
3236        Self::from_parts(EagerTensorParts {
3237            ctx,
3238            key,
3239            requires_grad,
3240            trace,
3241            semantic_trace,
3242            value,
3243            register_value: false,
3244        })
3245    }
3246
3247    pub(crate) fn new_result_value(
3248        ctx: Arc<EagerRuntime>,
3249        key: ValueKey<StdTensorOp>,
3250        value: TensorValue,
3251        requires_grad: bool,
3252        trace: Option<EagerTrace>,
3253        semantic_trace: Option<TracedTensor>,
3254    ) -> Result<Self> {
3255        let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
3256            Error::runtime_state(
3257                "EagerTensor::new_result_value",
3258                ErrorPhase::Execution,
3259                "a TensorValue could not be transferred into its allocation group",
3260            )
3261        })?;
3262        let value = AdValueRecord::from_group(group, slot, dtype, shape);
3263        Self::from_parts(EagerTensorParts {
3264            ctx,
3265            key,
3266            requires_grad,
3267            trace,
3268            semantic_trace,
3269            value,
3270            register_value: true,
3271        })
3272    }
3273
3274    fn from_parts(parts: EagerTensorParts) -> Result<Self> {
3275        let EagerTensorParts {
3276            ctx,
3277            key,
3278            requires_grad,
3279            trace,
3280            semantic_trace,
3281            value,
3282            register_value,
3283        } = parts;
3284        let grad_slot = Arc::new(Mutex::new(None));
3285        if requires_grad {
3286            ctx.try_register_grad_slot(&key, &grad_slot)?;
3287        }
3288        let record = Arc::new(EagerTensorRecord {
3289            value: Arc::clone(&value),
3290            key: key.clone(),
3291            trace: trace.clone(),
3292            semantic_trace: semantic_trace.clone(),
3293            requires_grad,
3294            grad_slot: Arc::clone(&grad_slot),
3295            ctx: Arc::clone(&ctx),
3296        });
3297        if register_value {
3298            ctx.try_register_value_record(&key, &record)?;
3299        }
3300
3301        Ok(Self {
3302            key,
3303            trace,
3304            semantic_trace,
3305            requires_grad,
3306            grad_slot,
3307            ctx,
3308            _record: record,
3309        })
3310    }
3311
3312    pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
3313        let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_untracked_result")?;
3314        Ok(Self::new_untracked_value_record(ctx, value, None))
3315    }
3316
3317    pub(crate) fn new_untracked_value_result(
3318        ctx: Arc<EagerRuntime>,
3319        value: TensorValue,
3320    ) -> Result<Self> {
3321        Self::new_untracked_value_result_with_semantic_trace(ctx, value, None)
3322    }
3323
3324    pub(crate) fn new_untracked_value_result_with_semantic_trace(
3325        ctx: Arc<EagerRuntime>,
3326        value: TensorValue,
3327        semantic_trace: Option<TracedTensor>,
3328    ) -> Result<Self> {
3329        let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
3330            Error::runtime_state(
3331                "EagerTensor::new_untracked_value_result",
3332                ErrorPhase::Execution,
3333                "a TensorValue could not be transferred into its allocation group",
3334            )
3335        })?;
3336        let value = AdValueRecord::from_group(group, slot, dtype, shape);
3337        Ok(Self::new_untracked_value_record(ctx, value, semantic_trace))
3338    }
3339
3340    fn new_untracked_value_record(
3341        ctx: Arc<EagerRuntime>,
3342        value: Arc<AdValueRecord>,
3343        semantic_trace: Option<TracedTensor>,
3344    ) -> Self {
3345        let key = eager_val_key();
3346        let grad_slot = Arc::new(Mutex::new(None));
3347        let record = Arc::new(EagerTensorRecord {
3348            value,
3349            key: key.clone(),
3350            trace: None,
3351            semantic_trace: semantic_trace.clone(),
3352            requires_grad: false,
3353            grad_slot: Arc::clone(&grad_slot),
3354            ctx: Arc::clone(&ctx),
3355        });
3356        Self {
3357            key,
3358            trace: None,
3359            semantic_trace,
3360            requires_grad: false,
3361            grad_slot,
3362            ctx,
3363            _record: record,
3364        }
3365    }
3366
3367    pub(crate) fn from_record(record: Arc<EagerTensorRecord>) -> Self {
3368        Self {
3369            key: record.key.clone(),
3370            trace: record.trace.clone(),
3371            semantic_trace: record.semantic_trace.clone(),
3372            requires_grad: record.requires_grad,
3373            grad_slot: Arc::clone(&record.grad_slot),
3374            ctx: Arc::clone(&record.ctx),
3375            _record: record,
3376        }
3377    }
3378
3379    /// Detach this tensor from the reverse graph.
3380    ///
3381    /// The returned tensor keeps the concrete value but no longer contributes
3382    /// gradients to the original graph.
3383    ///
3384    /// # Examples
3385    ///
3386    /// ```
3387    /// use tenferro_cpu::CpuBackend;
3388    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3389    ///
3390    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3391    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
3392    /// let y = x.detach();
3393    ///
3394    /// assert_eq!(y.value()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
3395    /// assert!(y.grad().unwrap().is_none());
3396    /// # Ok::<(), tenferro_ad::Error>(())
3397    /// ```
3398    pub fn detach(&self) -> Self {
3399        let semantic_trace = self
3400            .duplicate_value()
3401            .ok()
3402            .and_then(|tensor| TracedTensor::from_tensor_symbolic_shape(tensor).ok());
3403        Self::new_untracked_value_record(
3404            self.ctx.clone(),
3405            Arc::clone(&self._record.value),
3406            semantic_trace,
3407        )
3408    }
3409
3410    /// Detach this tensor from its graph and re-register it in a different
3411    /// context as an untracked leaf.
3412    ///
3413    /// # Examples
3414    ///
3415    /// ```
3416    /// use tenferro_cpu::CpuBackend;
3417    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3418    ///
3419    /// let ctx_a = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3420    /// let ctx_b = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3421    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx_a)?;
3422    /// let d = x.detach_into(&ctx_b)?;
3423    ///
3424    /// assert!(!d.tracks_grad());
3425    /// assert_eq!(d.ctx_id(), ctx_b.id());
3426    /// # Ok::<(), tenferro_ad::Error>(())
3427    /// ```
3428    ///
3429    /// # Errors
3430    ///
3431    /// Returns [`Error::RuntimeState`] if the source cannot be materialized or
3432    /// the target context cannot register its metadata.
3433    pub fn detach_into(&self, ctx: &Arc<EagerRuntime>) -> Result<Self> {
3434        Self::from_tensor_in(self.to_tensor()?, Arc::clone(ctx))
3435    }
3436
3437    /// Borrow the retained value without creating an owner or copy.
3438    ///
3439    /// # Errors
3440    ///
3441    /// Returns [`Error::RuntimeState`] when the retained allocation-group
3442    /// descriptor is unavailable or invalid.
3443    pub fn value(&self) -> Result<ValueGuard<'_>> {
3444        self._record.value.value("EagerTensor::value")
3445    }
3446
3447    /// Explicitly duplicate this value into a fresh standalone allocation.
3448    ///
3449    /// # Errors
3450    ///
3451    /// Returns [`Error::RuntimeState`] when the retained value or execution
3452    /// session is unavailable, or a typed host/backend error when the value
3453    /// cannot be materialized as a contiguous tensor.
3454    ///
3455    /// # Examples
3456    ///
3457    /// ```
3458    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3459    /// use tenferro_cpu::CpuBackend;
3460    ///
3461    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3462    /// let value = EagerTensor::from_tensor_in(
3463    ///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?,
3464    ///     ctx,
3465    /// )?;
3466    /// let duplicate = value.duplicate_value()?;
3467    /// assert_eq!(duplicate.as_slice::<f64>()?, &[1.0, 2.0]);
3468    /// # Ok::<(), tenferro_ad::Error>(())
3469    /// ```
3470    pub fn duplicate_value(&self) -> Result<Tensor> {
3471        let value = self.value()?;
3472        match value.duplicate_host_tensor() {
3473            Ok(tensor) => Ok(tensor),
3474            Err(_) => {
3475                let read = self
3476                    ._record
3477                    .value
3478                    .tensor_read("EagerTensor::duplicate_value")?;
3479                self.ctx
3480                    .with_execution_session(|session| session.to_contiguous_read(read))?
3481                    .map_err(Error::from)
3482            }
3483        }
3484    }
3485
3486    // INVARIANT: the error variants return the unchanged eager handle so a
3487    // caller can retry ownership extraction without an implicit copy.
3488    #[allow(clippy::result_large_err)]
3489    /// Consume this handle and structurally extract its retained allocation.
3490    ///
3491    /// A shared handle is returned unchanged as [`IntoValueError::NotUnique`].
3492    /// Group extraction failures return the unchanged handle and typed group
3493    /// error; no copy or fallback materialization is attempted.
3494    ///
3495    /// # Errors
3496    ///
3497    /// Returns [`IntoValueError::NotUnique`] when another handle retains the
3498    /// value, or [`IntoValueError::Extract`] when structural group extraction
3499    /// fails because the allocation is aliased or its descriptor is invalid.
3500    ///
3501    /// # Examples
3502    ///
3503    /// ```
3504    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3505    /// use tenferro_cpu::CpuBackend;
3506    ///
3507    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3508    /// let value = EagerTensor::from_tensor_in(
3509    ///     Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?,
3510    ///     ctx,
3511    /// )?;
3512    /// let owner = value
3513    ///     .into_value()
3514    ///     .expect("a uniquely owned value should be extractable");
3515    /// assert_eq!(owner.as_slice::<f64>()?, &[3.0]);
3516    /// # Ok::<(), tenferro_ad::Error>(())
3517    /// ```
3518    pub fn into_value(self) -> std::result::Result<Tensor, IntoValueError<Self>> {
3519        if Arc::strong_count(&self._record) != 1 {
3520            return Err(IntoValueError::NotUnique(self));
3521        }
3522        let Self { _record, .. } = self;
3523        let record = match Arc::try_unwrap(_record) {
3524            Ok(record) => record,
3525            Err(record) => return Err(IntoValueError::NotUnique(Self::from_record(record))),
3526        };
3527        let EagerTensorRecord {
3528            value,
3529            key,
3530            trace,
3531            semantic_trace,
3532            requires_grad,
3533            grad_slot,
3534            ctx,
3535        } = record;
3536        let value = match Arc::try_unwrap(value) {
3537            Ok(value) => value,
3538            Err(value) => {
3539                let record = Arc::new(EagerTensorRecord {
3540                    value,
3541                    key,
3542                    trace,
3543                    semantic_trace,
3544                    requires_grad,
3545                    grad_slot,
3546                    ctx,
3547                });
3548                return Err(IntoValueError::NotUnique(Self::from_record(record)));
3549            }
3550        };
3551        let AdValueRecord {
3552            container,
3553            slot,
3554            dtype,
3555            shape,
3556        } = value;
3557        let container = match Arc::try_unwrap(container) {
3558            Ok(container) => container,
3559            Err(container) => {
3560                let record = Arc::new(EagerTensorRecord {
3561                    value: Arc::new(AdValueRecord {
3562                        container,
3563                        slot,
3564                        dtype,
3565                        shape,
3566                    }),
3567                    key,
3568                    trace,
3569                    semantic_trace,
3570                    requires_grad,
3571                    grad_slot,
3572                    ctx,
3573                });
3574                return Err(IntoValueError::NotUnique(Self::from_record(record)));
3575            }
3576        };
3577        match container.group.into_tensor(slot) {
3578            Ok(tensor) => Ok(tensor),
3579            Err((group, error)) => {
3580                let record = Arc::new(EagerTensorRecord {
3581                    value: Arc::new(AdValueRecord {
3582                        container: Arc::new(RetentionContainer { group }),
3583                        slot,
3584                        dtype,
3585                        shape,
3586                    }),
3587                    key,
3588                    trace,
3589                    semantic_trace,
3590                    requires_grad,
3591                    grad_slot,
3592                    ctx,
3593                });
3594                Err(IntoValueError::Extract {
3595                    value: Self::from_record(record),
3596                    error,
3597                })
3598            }
3599        }
3600    }
3601
3602    /// Return this tensor's scalar dtype without materializing through
3603    /// [`value`](Self::value).
3604    pub fn dtype(&self) -> DType {
3605        self._record.value.dtype()
3606    }
3607
3608    /// Return this tensor's logical shape without materializing through
3609    /// [`value`](Self::value).
3610    pub fn shape(&self) -> &[usize] {
3611        self._record.value.shape()
3612    }
3613
3614    /// Borrow this tensor value as a [`TensorRead`].
3615    ///
3616    /// This is the preferred borrowed input boundary for executor calls. It
3617    /// preserves the option to replace eager storage with non-contiguous views
3618    /// without forcing callers through [`value`](Self::value).
3619    ///
3620    /// # Panics
3621    ///
3622    /// Panics if a validated eager value record becomes unavailable, which
3623    /// indicates an internal invariant violation.
3624    pub fn tensor_read(&self) -> TensorRead<'_> {
3625        self._record
3626            .value
3627            .tensor_read("EagerTensor::tensor_read")
3628            .expect("validated eager value record")
3629    }
3630
3631    /// Materialize this eager tensor as an owned [`Tensor`].
3632    ///
3633    /// This is the owned materialization boundary for callers that need a
3634    /// standalone compact tensor. The operation is fallible because eager
3635    /// values may be backed by lazy or backend-resident storage.
3636    ///
3637    /// # Errors
3638    ///
3639    /// Returns [`Error::RuntimeState`] if backend state is unavailable, or a
3640    /// typed tensor backend error when contiguous materialization fails.
3641    pub fn to_tensor(&self) -> Result<Tensor> {
3642        self.duplicate_value()
3643    }
3644
3645    /// Return the accumulated gradient currently stored for this tensor.
3646    ///
3647    /// The stored gradient accumulates across repeated `backward()` calls
3648    /// until it is cleared explicitly.
3649    ///
3650    /// For complex scalar losses, stored gradients use tenferro's
3651    /// Hermitian-adjoint cotangent convention. See
3652    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
3653    ///
3654    /// # Examples
3655    ///
3656    /// ```
3657    /// use tenferro_cpu::CpuBackend;
3658    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3659    ///
3660    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3661    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx).unwrap();
3662    /// let loss = x.exp().unwrap().reduce_sum(Some(&[0])).unwrap();
3663    /// let _cotangents = loss.backward().unwrap();
3664    ///
3665    /// let grad = x.grad()?.unwrap();
3666    /// assert_eq!(grad.shape(), &[2]);
3667    /// # Ok::<(), tenferro_ad::Error>(())
3668    /// ```
3669    ///
3670    /// # Errors
3671    ///
3672    /// Returns [`Error::RuntimeState`] if the gradient slot is poisoned or no
3673    /// longer available.
3674    pub fn grad(&self) -> Result<Option<GradientValue>> {
3675        self.grad_slot
3676            .lock()
3677            .map_err(|_| {
3678                Error::runtime_state(
3679                    "eager_gradient_slot",
3680                    ErrorPhase::Execution,
3681                    "lock poisoned",
3682                )
3683            })
3684            .map(|slot| {
3685                slot.as_ref().map(|record| GradientValue {
3686                    record: Arc::clone(record),
3687                    ctx: Arc::clone(&self.ctx),
3688                })
3689            })
3690    }
3691
3692    /// Clear the accumulated gradient stored for this tensor.
3693    ///
3694    /// This only affects this tensor's gradient slot. Other tensors in the
3695    /// same context retain their gradients until they are cleared explicitly or
3696    /// overwritten by later accumulation.
3697    ///
3698    /// # Examples
3699    ///
3700    /// ```
3701    /// use tenferro_cpu::CpuBackend;
3702    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3703    ///
3704    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3705    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx.clone()).unwrap();
3706    /// let y = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![4.0_f64, 5.0, 6.0]).unwrap(), ctx).unwrap();
3707    /// let loss = x.mul(&y).unwrap().reduce_sum(Some(&[0])).unwrap();
3708    /// let _ = loss.backward().unwrap();
3709    ///
3710    /// x.clear_grad()?;
3711    ///
3712    /// assert!(x.grad()?.is_none());
3713    /// assert!(y.grad()?.is_some());
3714    /// # Ok::<(), tenferro_ad::Error>(())
3715    /// ```
3716    ///
3717    /// # Errors
3718    ///
3719    /// Returns [`Error::RuntimeState`] if the gradient slot lock is poisoned.
3720    pub fn clear_grad(&self) -> Result<()> {
3721        *self.grad_slot.lock().map_err(|_| {
3722            Error::runtime_state(
3723                "eager_gradient_slot",
3724                ErrorPhase::Execution,
3725                "lock poisoned",
3726            )
3727        })? = None;
3728        Ok(())
3729    }
3730
3731    /// Report whether this tensor participates in gradient tracking.
3732    ///
3733    /// Tracked tensors keep a gradient slot in their eager context; untracked
3734    /// tensors and detached tensors do not.
3735    ///
3736    /// # Examples
3737    ///
3738    /// ```
3739    /// use tenferro_cpu::CpuBackend;
3740    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3741    ///
3742    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3743    /// let plain = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
3744    /// let tracked = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
3745    /// let detached = tracked.detach();
3746    ///
3747    /// assert!(!plain.tracks_grad());
3748    /// assert!(tracked.tracks_grad());
3749    /// assert!(!detached.tracks_grad());
3750    /// # Ok::<(), tenferro_ad::Error>(())
3751    /// ```
3752    pub fn tracks_grad(&self) -> bool {
3753        self.requires_grad
3754    }
3755
3756    #[cfg(test)]
3757    fn debug_trace_saved_value_count(&self) -> Option<usize> {
3758        None
3759    }
3760
3761    /// Return the opaque identifier of the context this tensor belongs to.
3762    ///
3763    /// # Examples
3764    ///
3765    /// ```
3766    /// use tenferro_cpu::CpuBackend;
3767    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3768    ///
3769    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3770    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
3771    ///
3772    /// assert_eq!(x.ctx_id(), ctx.id());
3773    /// # Ok::<(), tenferro_ad::Error>(())
3774    /// ```
3775    pub fn ctx_id(&self) -> ContextId {
3776        self.ctx.id()
3777    }
3778
3779    /// Borrow the eager runtime context that owns this tensor.
3780    pub fn runtime(&self) -> &Arc<EagerRuntime> {
3781        &self.ctx
3782    }
3783
3784    /// Check whether two tensors belong to the same eager context.
3785    ///
3786    /// # Examples
3787    ///
3788    /// ```
3789    /// use tenferro_cpu::CpuBackend;
3790    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3791    ///
3792    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3793    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
3794    /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx).unwrap();
3795    ///
3796    /// assert!(x.same_context(&y));
3797    /// # Ok::<(), tenferro_ad::Error>(())
3798    /// ```
3799    pub fn same_context(&self, other: &Self) -> bool {
3800        self.ctx_id() == other.ctx_id()
3801    }
3802
3803    #[cfg(test)]
3804    pub(crate) fn standard_graph_op(
3805        inputs: &[&Self],
3806        build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
3807    ) -> Result<Vec<Self>> {
3808        let Some(first) = inputs.first() else {
3809            return Err(Error::Internal(
3810                "standard eager graph op requires at least one input tensor".to_string(),
3811            ));
3812        };
3813        let ctx = Arc::clone(&first.ctx);
3814        for tensor in inputs.iter().skip(1) {
3815            if !first.same_context(tensor) {
3816                return Err(Error::ContextMismatch {
3817                    lhs: first.ctx_id(),
3818                    rhs: tensor.ctx_id(),
3819                });
3820            }
3821        }
3822
3823        let graph_input_keys = (0..inputs.len())
3824            .map(|_| next_input_key())
3825            .collect::<Vec<_>>();
3826        let graph = build_graph(&graph_input_keys)?;
3827        let initial_data = graph_input_keys
3828            .iter()
3829            .zip(inputs.iter())
3830            .map(|(key, tensor)| Ok((ValueKey::Input(key.clone()), tensor.to_tensor()?)))
3831            .collect::<Result<HashMap<_, _>>>()?;
3832        let execution = ctx.exec_standard_graph_outputs(graph.as_ref(), initial_data)?;
3833        if execution.outputs.len() != graph.outputs().len() {
3834            return Err(Error::Internal(format!(
3835                "standard eager graph op expected {} graph outputs, got {}",
3836                graph.outputs().len(),
3837                execution.outputs.len()
3838            )));
3839        }
3840
3841        if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
3842            return execution
3843                .outputs
3844                .into_iter()
3845                .map(|output| {
3846                    Self::new_unregistered_result_with_semantic_trace(
3847                        Arc::clone(&ctx),
3848                        eager_val_key(),
3849                        output,
3850                        false,
3851                        None,
3852                        None,
3853                    )
3854                })
3855                .collect();
3856        }
3857
3858        let recorded = record_eager_graph_outputs(
3859            graph.as_ref(),
3860            &graph_input_keys,
3861            &execution.outputs,
3862            inputs,
3863        )?;
3864        if recorded.traces.len() != execution.outputs.len() {
3865            return Err(Error::Internal(format!(
3866                "standard eager graph op expected {} eager traces, got {}",
3867                execution.outputs.len(),
3868                recorded.traces.len()
3869            )));
3870        }
3871
3872        recorded
3873            .traces
3874            .into_iter()
3875            .zip(recorded.semantic_traces)
3876            .zip(execution.outputs)
3877            .map(|((trace, semantic_trace), output)| {
3878                Self::new_result_with_semantic_trace(
3879                    Arc::clone(&ctx),
3880                    trace.key,
3881                    output,
3882                    trace.requires_grad,
3883                    trace.trace,
3884                    semantic_trace,
3885                )
3886            })
3887            .collect()
3888    }
3889
3890    /// Run reverse-mode AD from this scalar output.
3891    ///
3892    /// Returns the full cotangent map produced by the reverse pass and also
3893    /// accumulates into `grad()` for tracked eager tensors reachable from this
3894    /// output.
3895    ///
3896    /// For complex scalar outputs, cotangents use tenferro's Hermitian
3897    /// real-inner-product convention. See
3898    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
3899    ///
3900    /// # Examples
3901    ///
3902    /// ```
3903    /// use tenferro_cpu::CpuBackend;
3904    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3905    ///
3906    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3907    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx).unwrap();
3908    /// let loss = x.add(&x).unwrap().reduce_sum(Some(&[0])).unwrap();
3909    /// let _cotangents = loss.backward().unwrap();
3910    /// let loss = x.add(&x).unwrap().reduce_sum(Some(&[0])).unwrap();
3911    /// let _cotangents = loss.backward().unwrap();
3912    ///
3913    /// assert_eq!(x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(), &[4.0, 4.0, 4.0]);
3914    /// # Ok::<(), tenferro_ad::Error>(())
3915    /// ```
3916    ///
3917    /// # Errors
3918    ///
3919    /// Returns [`Error::NonScalarGrad`] when this output is not scalar,
3920    /// [`Error::UnsupportedAdRule`] when a graph operation lacks a reverse rule,
3921    /// or a typed validation/backend/runtime-state error during the reverse pass.
3922    pub fn backward(&self) -> Result<Gradients> {
3923        if !self.shape().is_empty() {
3924            return Err(Error::NonScalarGrad {
3925                shape: self.shape().to_vec(),
3926            });
3927        }
3928
3929        let value = self.to_tensor()?;
3930        let seed = {
3931            let mut backend = self.ctx.lock_backend()?;
3932            one_like_tensor(&value, &mut *backend)?
3933        };
3934        self.backward_from_seed(seed)
3935    }
3936
3937    /// Run reverse-mode AD from this output with an explicit cotangent seed.
3938    ///
3939    /// This is the stateful eager VJP sugar: it returns the cotangent map and
3940    /// accumulates reachable tracked leaves into their `grad()` slots. Use
3941    /// [`EagerRuntime::vjp`] when the VJP result should be returned as a
3942    /// composable eager tensor without touching grad slots.
3943    ///
3944    /// # Examples
3945    ///
3946    /// ```
3947    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3948    /// use tenferro_cpu::CpuBackend;
3949    ///
3950    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
3951    /// let x = EagerTensor::requires_grad_in(
3952    ///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
3953    ///     ctx.clone(),
3954    /// )?;
3955    /// let seed = EagerTensor::from_tensor_in(
3956    ///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
3957    ///     ctx,
3958    /// )?;
3959    /// let y = x.mul(&x)?;
3960    /// y.backward_with(&seed)?;
3961    /// assert_eq!(x.grad()?.unwrap().as_slice::<f64>().unwrap(), &[4.0, 12.0]);
3962    /// # Ok::<(), tenferro_ad::Error>(())
3963    /// ```
3964    ///
3965    /// # Errors
3966    ///
3967    /// Returns [`Error::ContextMismatch`] when `cotangent` belongs to another
3968    /// eager runtime, [`Error::Validation`] when its shape or dtype is not a
3969    /// valid seed, [`Error::UnsupportedAdRule`] for an unavailable reverse
3970    /// rule, or a typed backend/runtime-state error during execution.
3971    pub fn backward_with(&self, cotangent: &EagerTensor) -> Result<Gradients> {
3972        if !self.same_context(cotangent) {
3973            return Err(Error::ContextMismatch {
3974                lhs: self.ctx_id(),
3975                rhs: cotangent.ctx_id(),
3976            });
3977        }
3978        validate_seed_tensor("backward", self, cotangent)?;
3979        self.backward_from_seed(cotangent.to_tensor()?)
3980    }
3981
3982    fn backward_from_seed(&self, seed: Tensor) -> Result<Gradients> {
3983        let cotangent =
3984            EagerTensor::new_result(Arc::clone(&self.ctx), eager_val_key(), seed, false, None)?;
3985        let candidate_keys = {
3986            let mut slots = self.ctx.lock_grad_slots()?;
3987            let mut keys = Vec::new();
3988            slots.retain(|key, slot| {
3989                if slot.upgrade().is_some() {
3990                    keys.push(key.clone());
3991                    true
3992                } else {
3993                    false
3994                }
3995            });
3996            keys
3997        };
3998
3999        let mut cotangents = HashMap::new();
4000        for key in candidate_keys {
4001            let Some(record) = self.ctx.value_record(&key)? else {
4002                continue;
4003            };
4004            if !record.requires_grad {
4005                continue;
4006            }
4007            let wrt = EagerTensor::from_record(record);
4008            let Some(grad) = self.ctx.vjp_optional(self, &wrt, &cotangent)? else {
4009                continue;
4010            };
4011            let tensor = match grad.into_value() {
4012                Ok(tensor) => tensor,
4013                Err(IntoValueError::NotUnique(handle)) => handle.duplicate_value()?,
4014                Err(IntoValueError::Extract { error, .. }) => {
4015                    return Err(Error::runtime_state_source(
4016                        "EagerTensor::backward",
4017                        ErrorPhase::Execution,
4018                        error,
4019                    ));
4020                }
4021            };
4022            cotangents.insert(key, tensor);
4023        }
4024        let mut backend = self.ctx.lock_backend()?;
4025        self.ctx.store_grads(&cotangents, &mut backend)?;
4026        Gradients::from_tensors(cotangents)
4027    }
4028}
4029
4030pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
4031    ValueKey::Input(next_input_key())
4032}
4033
4034pub(crate) struct RecordedEagerTrace {
4035    pub(crate) key: ValueKey<StdTensorOp>,
4036    pub(crate) trace: Option<EagerTrace>,
4037    pub(crate) requires_grad: bool,
4038}
4039
4040pub(crate) struct RecordedEagerOutputs {
4041    pub(crate) traces: Vec<RecordedEagerTrace>,
4042    pub(crate) semantic_traces: Vec<Option<TracedTensor>>,
4043}
4044
4045pub(crate) fn record_eager_outputs(
4046    op: &StdTensorOp,
4047    outputs: &[&Tensor],
4048    inputs: &[&EagerTensor],
4049) -> Result<RecordedEagerOutputs> {
4050    let output_metadata = outputs
4051        .iter()
4052        .map(|output| tensor_meta_from_tensor(output))
4053        .collect::<Vec<_>>();
4054    record_eager_outputs_inner(op, output_metadata, inputs)
4055}
4056
4057pub(crate) fn record_eager_value_outputs(
4058    op: &StdTensorOp,
4059    outputs: &[&TensorValue],
4060    inputs: &[&EagerTensor],
4061) -> Result<RecordedEagerOutputs> {
4062    let output_metadata = outputs
4063        .iter()
4064        .map(|output| tensor_meta_from_value(output))
4065        .collect::<Vec<_>>();
4066    record_eager_outputs_inner(op, output_metadata, inputs)
4067}
4068
4069fn record_eager_outputs_inner(
4070    op: &StdTensorOp,
4071    output_metadata: Vec<TensorMeta>,
4072    inputs: &[&EagerTensor],
4073) -> Result<RecordedEagerOutputs> {
4074    let semantic_traces = record_semantic_eager_outputs(op, &output_metadata, inputs)?;
4075    record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
4076}
4077
4078fn record_semantic_eager_outputs(
4079    op: &StdTensorOp,
4080    output_metadata: &[TensorMeta],
4081    inputs: &[&EagerTensor],
4082) -> Result<Vec<Option<TracedTensor>>> {
4083    // Materialize a constant semantic leaf for any untracked input that lost
4084    // its implicit semantic trace on the active-edge fast path. This keeps
4085    // "untracked constant feeds tracked AD" working (PyTorch-style: untracked
4086    // = constant leaf, no gradient flows to it) without re-recording every
4087    // untracked op at creation time.
4088    let mut owned_constants = Vec::<TracedTensor>::new();
4089    for input in inputs {
4090        if input.semantic_trace.is_none() {
4091            owned_constants.push(TracedTensor::from_tensor_symbolic_shape(
4092                input.to_tensor()?,
4093            )?);
4094        }
4095    }
4096    let mut constants = owned_constants.iter();
4097    let mut semantic_inputs: Vec<&TracedTensor> = inputs
4098        .iter()
4099        .map(|input| {
4100            input
4101                .semantic_trace
4102                .as_ref()
4103                .unwrap_or_else(|| constants.next().expect("materialized constant"))
4104        })
4105        .collect();
4106    let promotion_plan =
4107        eager_input_promotion_plan(op, inputs.len(), |index| inputs[index].dtype());
4108    // Mirror eager execution in the deferred carrier only. The concrete
4109    // tensors have already been promoted at the execution boundary, so these
4110    // casts add semantic graph nodes without an eager copy or backend kernel.
4111    let promoted_semantic_inputs = if semantic_inputs.iter().enumerate().any(|(index, semantic)| {
4112        semantic.dtype != promotion_plan.target_dtype(index, semantic.dtype)
4113    }) {
4114        Some(
4115            semantic_inputs
4116                .iter()
4117                .enumerate()
4118                .map(|(index, &semantic)| {
4119                    let target = promotion_plan.target_dtype(index, semantic.dtype);
4120                    if semantic.dtype == target {
4121                        Ok(Cow::Borrowed(semantic))
4122                    } else {
4123                        semantic.cast(target).map(Cow::Owned)
4124                    }
4125                })
4126                .collect::<Result<Vec<Cow<'_, TracedTensor>>>>()?,
4127        )
4128    } else {
4129        None
4130    };
4131    if let Some(promoted_semantic_inputs) = &promoted_semantic_inputs {
4132        semantic_inputs = promoted_semantic_inputs.iter().map(Cow::as_ref).collect();
4133    }
4134    let exact_semantic_inputs = if matches!(op, StdTensorOp::Concatenate { .. }) {
4135        Some(
4136            semantic_inputs
4137                .iter()
4138                .zip(inputs)
4139                .map(|(&semantic, input)| {
4140                    if semantic.is_concrete_shape() {
4141                        Ok(semantic.clone())
4142                    } else {
4143                        semantic.reshape(input.shape())
4144                    }
4145                })
4146                .collect::<Result<Vec<_>>>()?,
4147        )
4148    } else {
4149        None
4150    };
4151    if let Some(exact_semantic_inputs) = &exact_semantic_inputs {
4152        semantic_inputs = exact_semantic_inputs.iter().collect();
4153    }
4154    // Deferred materialization (issue #1665 steps 6-7): append only a raw
4155    // carrier. The runtime helper retains metadata scopes introduced by the
4156    // promotion/exactification helpers without analyzing this operation.
4157    let outputs = tenferro_runtime::extension::append_raw_eager_outputs(
4158        op.clone(),
4159        &semantic_inputs,
4160        output_metadata,
4161    )?;
4162    Ok(outputs.into_iter().map(Some).collect())
4163}
4164
4165#[cfg(test)]
4166fn record_eager_graph_outputs(
4167    graph: &Graph<StdTensorOp>,
4168    graph_input_keys: &[TensorInputKey],
4169    outputs: &[Tensor],
4170    inputs: &[&EagerTensor],
4171) -> Result<RecordedEagerOutputs> {
4172    let semantic_traces = record_semantic_eager_graph_outputs(graph, graph_input_keys, inputs)?;
4173    let output_metadata = outputs.iter().map(tensor_meta_from_tensor);
4174    record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
4175}
4176
4177#[cfg(test)]
4178fn record_semantic_eager_graph_outputs(
4179    graph: &Graph<StdTensorOp>,
4180    graph_input_keys: &[TensorInputKey],
4181    inputs: &[&EagerTensor],
4182) -> Result<Vec<Option<TracedTensor>>> {
4183    let Some(semantic_inputs) = inputs
4184        .iter()
4185        .map(|input| input.semantic_trace.as_ref())
4186        .collect::<Option<Vec<_>>>()
4187    else {
4188        return Ok(vec![None; graph.outputs().len()]);
4189    };
4190    if graph_input_keys.len() != semantic_inputs.len() {
4191        return Err(Error::Internal(format!(
4192            "semantic graph recording expected {} input keys, got {}",
4193            semantic_inputs.len(),
4194            graph_input_keys.len()
4195        )));
4196    }
4197
4198    let mut values = HashMap::new();
4199    for (key, tensor) in graph_input_keys.iter().zip(semantic_inputs) {
4200        values.insert(ValueKey::Input(key.clone()), tensor.clone());
4201    }
4202
4203    for op_node in graph.operations() {
4204        let input_values = op_node
4205            .inputs
4206            .iter()
4207            .map(|input| {
4208                let key = match input {
4209                    ValueRef::Local(local_id) => &graph.values()[*local_id].key,
4210                    ValueRef::External(key) => key,
4211                };
4212                values.get(key).cloned().ok_or_else(|| {
4213                    Error::Internal(format!(
4214                        "semantic graph recording missing value for {key:?}"
4215                    ))
4216                })
4217            })
4218            .collect::<Result<Vec<_>>>()?;
4219        let input_refs = input_values.iter().collect::<Vec<_>>();
4220        let semantic_outputs = match &op_node.operation {
4221            StdTensorOp::Extension(ext) => {
4222                tenferro_runtime::extension::apply(Arc::clone(ext), &input_refs)?
4223            }
4224            op => tenferro_runtime::extension::apply_standard_op(op.clone(), &input_refs)?,
4225        };
4226        if semantic_outputs.len() != op_node.outputs.len() {
4227            return Err(Error::Internal(format!(
4228                "semantic graph recording expected {} outputs for {:?}, got {}",
4229                op_node.outputs.len(),
4230                op_node.operation,
4231                semantic_outputs.len()
4232            )));
4233        }
4234        for (output_id, output) in op_node.outputs.iter().copied().zip(semantic_outputs) {
4235            values.insert(graph.values()[output_id].key.clone(), output);
4236        }
4237    }
4238
4239    graph
4240        .outputs()
4241        .iter()
4242        .map(|&output_id| {
4243            let key = &graph.values()[output_id].key;
4244            values.get(key).cloned().map(Some).ok_or_else(|| {
4245                Error::Internal(format!(
4246                    "semantic graph recording missing output for {key:?}"
4247                ))
4248            })
4249        })
4250        .collect()
4251}
4252
4253fn record_eager_outputs_from_metadata(
4254    output_metadata: impl IntoIterator<Item = TensorMeta>,
4255    semantic_traces: Vec<Option<TracedTensor>>,
4256    inputs: &[&EagerTensor],
4257) -> Result<RecordedEagerOutputs> {
4258    let output_metadata = output_metadata.into_iter().collect::<Vec<_>>();
4259    if semantic_traces.len() != output_metadata.len() {
4260        return Err(Error::Internal(format!(
4261            "eager recording expected {} semantic traces, got {}",
4262            output_metadata.len(),
4263            semantic_traces.len()
4264        )));
4265    }
4266    let requires_grad =
4267        eager_grad_recording_enabled() && inputs.iter().any(|input| input.requires_grad);
4268    let trace_count = output_metadata.len();
4269    let traces = (0..trace_count)
4270        .map(|_| RecordedEagerTrace {
4271            key: eager_val_key(),
4272            trace: None,
4273            requires_grad,
4274        })
4275        .collect();
4276
4277    Ok(RecordedEagerOutputs {
4278        traces,
4279        semantic_traces,
4280    })
4281}
4282
4283fn tensor_meta_from_value(value: &TensorValue) -> TensorMeta {
4284    TensorMeta::exact(
4285        value.dtype(),
4286        value.shape().iter().copied().map(SymDim::from).collect(),
4287    )
4288}
4289
4290pub(crate) fn exec_single_output(
4291    op: &StdTensorOp,
4292    inputs: &[&Tensor],
4293    ctx: &EagerRuntime,
4294) -> Result<Tensor> {
4295    let mut outputs = ctx.exec_outputs(op, inputs)?;
4296    if outputs.len() != 1 {
4297        return Err(Error::Internal(format!(
4298            "expected one eager output for {:?}, got {}",
4299            op,
4300            outputs.len()
4301        )));
4302    }
4303    Ok(profile_eager_op_section(
4304        "exec_single_output.remove_output",
4305        || outputs.remove(0),
4306    ))
4307}
4308
4309pub(crate) fn exec_single_output_read(
4310    op: &StdTensorOp,
4311    inputs: &[TensorRead<'_>],
4312    ctx: &EagerRuntime,
4313) -> Result<Tensor> {
4314    let mut outputs = ctx.exec_outputs_read(op, inputs)?;
4315    if outputs.len() != 1 {
4316        return Err(Error::Internal(format!(
4317            "expected one eager output for {:?}, got {}",
4318            op,
4319            outputs.len()
4320        )));
4321    }
4322    Ok(profile_eager_op_section(
4323        "exec_single_output_read.remove_output",
4324        || outputs.remove(0),
4325    ))
4326}
4327
4328#[cfg(test)]
4329pub(crate) fn zero_like_tensor<B: TensorBackend>(
4330    input: &Tensor,
4331    backend: &mut B,
4332) -> Result<Tensor> {
4333    let host = match input {
4334        Tensor::F32(tensor) => Tensor::F32(TypedTensor::zeros(tensor.shape().to_vec())?),
4335        Tensor::F64(tensor) => Tensor::F64(TypedTensor::zeros(tensor.shape().to_vec())?),
4336        Tensor::I32(tensor) => Tensor::I32(TypedTensor::zeros(tensor.shape().to_vec())?),
4337        Tensor::I64(tensor) => Tensor::I64(TypedTensor::zeros(tensor.shape().to_vec())?),
4338        Tensor::Bool(tensor) => Tensor::Bool(TypedTensor::from_vec_col_major(
4339            tensor.shape().to_vec(),
4340            vec![false; tensor.n_elements()],
4341        )?),
4342        Tensor::C32(tensor) => Tensor::C32(TypedTensor::zeros(tensor.shape().to_vec())?),
4343        Tensor::C64(tensor) => Tensor::C64(TypedTensor::zeros(tensor.shape().to_vec())?),
4344    };
4345    backend
4346        .upload_host_tensor(TensorRead::from_tensor(&host))
4347        .map_err(Error::from)
4348}
4349
4350pub(crate) fn one_like_tensor<B: TensorBackend>(input: &Tensor, backend: &mut B) -> Result<Tensor> {
4351    let host = ones_tensor(input.dtype(), input.shape().to_vec())?;
4352    backend
4353        .upload_host_tensor(TensorRead::from_tensor(&host))
4354        .map_err(Error::from)
4355}
4356
4357#[cfg(test)]
4358mod tests;