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;
16use num_complex::Complex64;
17
18use crate::extension::{
19    validate_eager_extension_target, EagerExtensionBackendKind, EagerExtensionTarget,
20};
21
22#[cfg(test)]
23use computegraph::graph::Graph;
24use computegraph::ValueKey;
25#[cfg(test)]
26use computegraph::ValueRef;
27#[cfg(test)]
28use tenferro_ops::input_key::TensorInputKey;
29use tenferro_ops::{std_tensor_op::StdTensorOp, SymDim, TensorMeta};
30use tenferro_runtime::ad_support::{
31    analyze_deferred_semantic_trace, compile_ad_source, ones_tensor, RetainedValue,
32};
33use tenferro_runtime::program::{ProgramValueMetadata, SemanticFingerprint, SemanticProgram};
34use tenferro_runtime::{
35    CompiledGraph, EngineId, EngineTarget, ErrorPhase, ExecutionContextIdentity, ExecutionEngine,
36    ExtensionModule, GraphCompiler, Prepared, Runtime, RuntimeConfigError, RuntimeEpoch,
37    TracedTensor,
38};
39use tenferro_tensor::Engine;
40#[cfg(test)]
41use tenferro_tensor::TypedTensor;
42use tenferro_tensor::{
43    AllocationGroup, CacheStats, CompareDir, DType, DescriptorSlot, DotGeneralConfig, GatherConfig,
44    GroupError, IntoShapeVec, PadConfig, ScatterConfig, SliceConfig, Tensor, TensorRead,
45    TensorScalar, TensorValue, TensorView,
46};
47
48use crate::eager_backend::{
49    cpu_host_owned_read, engine_kind_registration, unregisterable_engine_error,
50};
51#[cfg(test)]
52use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
53use crate::eager_exec::{eager_input_promotion_plan, exec_extension_op_on_tensor_reads};
54use crate::error::{ContextId, Error, Result};
55use crate::metadata::tensor_meta_from_tensor;
56use crate::semantic_extension::SemanticExtensionRuleSet;
57use crate::traced::{derivative_trace_from_frozen_program, next_input_key};
58use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
59
60use crate::AdContext;
61
62pub(crate) type GradSlot = Arc<Mutex<Option<Arc<AdValueRecord>>>>;
63pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<AdValueRecord>>>>;
64
65mod composite;
66mod residuals;
67pub(crate) use residuals::{finish_residuals, EagerTrace};
68
69#[derive(Debug, Default, Clone)]
70struct EagerOpProfileEntry {
71    calls: usize,
72    total_time: Duration,
73}
74
75thread_local! {
76    static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
77        RefCell::new(HashMap::new());
78    static EAGER_NO_GRAD_DEPTH: Cell<usize> = const { Cell::new(0) };
79    static EAGER_CAPTURE_DEPTH: Cell<usize> = const { Cell::new(0) };
80    #[cfg(test)]
81    static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
82    #[cfg(test)]
83    static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
84    #[cfg(test)]
85    static EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
86}
87
88#[cfg(test)]
89pub(crate) static EAGER_SEMANTIC_VJP_EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
90
91pub(crate) fn eager_grad_recording_enabled() -> bool {
92    EAGER_NO_GRAD_DEPTH.with(|depth| depth.get() == 0)
93}
94
95pub(crate) fn eager_capture_active() -> bool {
96    EAGER_CAPTURE_DEPTH.with(|depth| depth.get() > 0)
97}
98
99fn eager_semantic_vjp_enabled() -> bool {
100    #[cfg(test)]
101    if let Some(value) = EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
102        return value;
103    }
104
105    // Semantic eager VJP/JVP on by default (Unification 7).
106    // Set TENFERRO_EAGER_SEMANTIC_VJP=0 to disable.
107    static ENABLED: OnceLock<bool> = OnceLock::new();
108    *ENABLED.get_or_init(|| env::var("TENFERRO_EAGER_SEMANTIC_VJP").map_or(true, |v| v != "0"))
109}
110
111/// Scope guard that temporarily disables eager operation recording.
112///
113/// Values computed while this guard is alive are concrete eager tensors, but
114/// they do not participate in reverse-mode gradient tracking.
115///
116/// # Examples
117///
118/// ```
119/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
120///
121/// let ctx = EagerRuntime::new()?;
122/// let x = EagerTensor::requires_grad_in(
123///     Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
124///     ctx.clone(),
125/// )?;
126/// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
127///     let _guard = ctx.no_grad();
128///     s.mul(&x, &x)
129/// })?;
130/// assert!(!y.tracks_grad());
131/// # Ok::<(), tenferro_ad::Error>(())
132/// ```
133#[derive(Debug)]
134pub struct EagerNoGradGuard {
135    active: bool,
136    // Thread-local depth guard: must not be Send so it cannot be moved to and
137    // dropped on another thread (which would corrupt the creator's depth).
138    _not_send: PhantomData<Rc<()>>,
139}
140
141impl Drop for EagerNoGradGuard {
142    fn drop(&mut self) {
143        if !self.active {
144            return;
145        }
146        EAGER_NO_GRAD_DEPTH.with(|depth| {
147            depth.set(depth.get().saturating_sub(1));
148        });
149        self.active = false;
150    }
151}
152
153/// Scope guard that keeps semantic-trace recording active for untracked
154/// intermediates.
155///
156/// Under active-edge semantics (issue #1665 Def 1), an operation whose inputs
157/// are all untracked produces no autograd nodes and drops its semantic trace.
158/// Inside this guard, such operations still record their semantic trace, so a
159/// later functional JVP/VJP can differentiate with respect to an untracked or
160/// detached leaf. This replaces the pre-Def-1 implicit recording.
161///
162/// # Examples
163///
164/// ```
165/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
166///
167/// let ctx = EagerRuntime::new()?;
168/// let x = EagerTensor::from_tensor_in(
169///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
170///     ctx.clone(),
171/// )?;
172/// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
173///     let _capture = ctx.capture_trace();
174///     s.mul(&x, &x)
175/// })?;
176/// let seed = EagerTensor::from_tensor_in(
177///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 1.0]).unwrap(),
178///     ctx.clone(),
179/// )?;
180/// let dx = ctx.vjp(&mut tenferro_cpu::CpuEngine::serial(), &y, &x, &seed)?;
181/// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[2.0, 4.0]);
182/// # Ok::<(), tenferro_ad::Error>(())
183/// ```
184#[derive(Debug)]
185pub struct EagerTraceCaptureGuard {
186    active: bool,
187    // Thread-local depth guard: must not be Send so it cannot be moved to and
188    // dropped on another thread (which would corrupt the creator's depth).
189    _not_send: PhantomData<Rc<()>>,
190}
191
192impl Drop for EagerTraceCaptureGuard {
193    fn drop(&mut self) {
194        if !self.active {
195            return;
196        }
197        EAGER_CAPTURE_DEPTH.with(|depth| {
198            depth.set(depth.get().saturating_sub(1));
199        });
200        self.active = false;
201    }
202}
203
204pub(crate) fn eager_op_profile_enabled() -> bool {
205    #[cfg(test)]
206    if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
207        return value;
208    }
209
210    static ENABLED: OnceLock<bool> = OnceLock::new();
211    *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
212}
213
214pub(crate) fn eager_op_profile_start() -> Option<Instant> {
215    eager_op_profile_enabled().then(Instant::now)
216}
217
218pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
219    if !eager_op_profile_enabled() {
220        return;
221    }
222    EAGER_OP_PROFILE_STATE.with(|state| {
223        let mut state = state.borrow_mut();
224        let entry = state.entry(section).or_default();
225        entry.calls += 1;
226        entry.total_time += elapsed;
227    });
228}
229
230pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
231    if !eager_op_profile_enabled() {
232        return f();
233    }
234    let started = Instant::now();
235    let result = f();
236    record_eager_op_profile(section, started.elapsed());
237    result
238}
239
240pub(crate) fn maybe_print_eager_op_profile() {
241    if !eager_op_profile_enabled() {
242        return;
243    }
244    let Some(print_every) = eager_op_profile_print_every() else {
245        return;
246    };
247    if print_every == 0 {
248        return;
249    }
250
251    let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
252        state
253            .borrow()
254            .get("nary_op.total")
255            .is_some_and(|entry| entry.calls % print_every == 0)
256    });
257    if should_print {
258        print_and_reset_eager_op_profile();
259    }
260}
261
262fn eager_op_profile_print_every() -> Option<usize> {
263    #[cfg(test)]
264    if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
265        return value;
266    }
267
268    env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
269        .ok()?
270        .parse()
271        .ok()
272}
273
274pub(crate) fn print_and_reset_eager_op_profile() {
275    EAGER_OP_PROFILE_STATE.with(|state| {
276        let mut entries: Vec<_> = state
277            .borrow()
278            .iter()
279            .map(|(section, entry)| (*section, entry.clone()))
280            .collect();
281        state.borrow_mut().clear();
282        entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
283
284        eprintln!("=== tenferro eager op profile ===");
285        for (section, entry) in entries {
286            let Some(per_call_us) = eager_op_profile_per_call_us(&entry) else {
287                continue;
288            };
289            eprintln!(
290                "{section}: calls={} total={:.6}ms per_call={:.3}us",
291                entry.calls,
292                entry.total_time.as_secs_f64() * 1.0e3,
293                per_call_us,
294            );
295        }
296    });
297}
298
299fn eager_op_profile_per_call_us(entry: &EagerOpProfileEntry) -> Option<f64> {
300    (entry.calls != 0).then(|| entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64)
301}
302
303fn runtime_config_error(op: &'static str, source: RuntimeConfigError) -> Error {
304    Error::runtime_state_source(op, ErrorPhase::Execution, source)
305}
306
307fn runtime_state_source<E>(op: &'static str, source: E) -> Error
308where
309    E: std::error::Error + Send + Sync + 'static,
310{
311    Error::runtime_state_source(op, ErrorPhase::Execution, source)
312}
313
314/// Run `prepared` on `engine` and return owned compact outputs.
315pub(crate) fn run_prepared_on(
316    runtime: &Runtime,
317    engine: &mut dyn ExecutionEngine,
318    prepared: &Prepared,
319    inputs: &[TensorRead<'_>],
320) -> Result<Vec<Tensor>> {
321    let outputs = runtime.run(&mut [&mut *engine], prepared, inputs)?;
322    outputs
323        .into_iter()
324        .map(|output| output.into_tensor(&mut *engine).map_err(Error::from))
325        .collect()
326}
327
328/// Run a session-capable prepared extension executor on `engine`, lending it
329/// the engine's extension caches for the call.
330pub(crate) fn execute_prepared_extension_on(
331    engine: &mut dyn ExecutionEngine,
332    executor: &dyn tenferro_runtime::PreparedOperationExecutor,
333    inputs: &[TensorRead<'_>],
334) -> Result<Vec<Tensor>> {
335    with_lent_extension_caches(engine, |engine, caches| {
336        executor.execute_in_session(engine, caches, inputs)
337    })
338}
339
340/// Lend `engine`'s extension caches to `callback` and restore them on every
341/// exit path, including unwind.
342///
343/// A callback panic also marks the engine broken before it is resumed: the
344/// caches may be half updated, so the engine must not accept more work.
345pub(crate) fn with_lent_extension_caches<R>(
346    engine: &mut dyn ExecutionEngine,
347    callback: impl FnOnce(&mut dyn ExecutionEngine, &mut tenferro_runtime::ExtensionCacheStore) -> R,
348) -> R {
349    let mut caches = engine.take_extension_caches();
350    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
351        callback(&mut *engine, &mut caches)
352    }));
353    engine.restore_extension_caches(caches);
354    match result {
355        Ok(value) => value,
356        Err(payload) => {
357            engine.mark_broken();
358            std::panic::resume_unwind(payload)
359        }
360    }
361}
362
363/// The eager extension backend kind of an engine target, or `None` for an
364/// engine type without eager extension support.
365pub(crate) fn eager_extension_backend_kind(
366    target: &EngineTarget,
367) -> Option<EagerExtensionBackendKind> {
368    if target.context_identity() == ExecutionContextIdentity::of::<tenferro_cpu::CpuEngine>() {
369        return Some(EagerExtensionBackendKind::Cpu);
370    }
371    #[cfg(feature = "cuda")]
372    if target.context_identity() == ExecutionContextIdentity::of::<tenferro_gpu::cuda::CudaEngine>()
373    {
374        return Some(EagerExtensionBackendKind::Cuda);
375    }
376    None
377}
378
379/// Stats for caches owned by an [`EagerRuntime`].
380///
381/// `retained_bytes` fields are logical payload estimates, not process RSS.
382#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
383pub struct EagerRuntimeCacheStats {
384    /// Eager AD transform memoization cache.
385    pub ad_transforms: CacheStats,
386    /// Prepared eager derivative program cache.
387    pub prepared_derivatives: CacheStats,
388}
389
390#[cfg(test)]
391pub(crate) struct EagerGraphExecution {
392    pub(crate) outputs: Vec<Tensor>,
393}
394
395/// A read-only value view retained by an eager tensor record.
396///
397/// The guard borrows the record's allocation group. It never owns a tensor and
398/// cannot be converted into a mutable view.
399///
400/// # Examples
401///
402/// ```
403/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
404///
405/// let ctx = EagerRuntime::new()?;
406/// let value = EagerTensor::from_tensor_in(
407///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?,
408///     ctx,
409/// )?;
410/// let view = value.value()?;
411/// assert_eq!(view.shape(), &[2]);
412/// # Ok::<(), tenferro_ad::Error>(())
413/// ```
414#[derive(Debug)]
415pub struct ValueGuard<'a> {
416    view: TensorView<'a>,
417}
418
419impl<'a> ValueGuard<'a> {
420    /// Return the scalar dtype of the retained value.
421    pub fn dtype(&self) -> DType {
422        self.view.dtype()
423    }
424
425    /// Return the logical shape of the retained value.
426    pub fn shape(&self) -> &[usize] {
427        self.view.shape()
428    }
429
430    /// Borrow the dtype-erased tensor view.
431    pub fn as_tensor_view(&self) -> &TensorView<'_> {
432        &self.view
433    }
434
435    /// Borrow compact host bytes through the tensor's explicit scalar type.
436    ///
437    /// Backend-resident values return the backend's typed host-access error;
438    /// this method does not download storage implicitly.
439    ///
440    /// # Errors
441    ///
442    /// Returns [`tenferro_tensor::ValidationError::DTypeMismatch`] when
443    /// `T` does not match the view dtype, [`tenferro_tensor::ValidationError::NonContiguousViewAsSlice`]
444    /// for a non-contiguous view, or [`tenferro_tensor::Error::HostAccess`]
445    /// when backend storage cannot be mapped as a host slice.
446    pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&'a [T]> {
447        self.view.as_slice()
448    }
449
450    fn duplicate_host_tensor(&self) -> tenferro_tensor::Result<Tensor> {
451        match &self.view {
452            TensorView::F32(view) => {
453                <f32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
454            }
455            TensorView::F64(view) => {
456                <f64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
457            }
458            TensorView::I32(view) => {
459                <i32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
460            }
461            TensorView::I64(view) => {
462                <i64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
463            }
464            TensorView::Bool(view) => <bool as TensorScalar>::into_tensor(
465                view.shape().to_vec(),
466                view.as_slice()?.to_vec(),
467            ),
468            TensorView::C32(view) => <num_complex::Complex32 as TensorScalar>::into_tensor(
469                view.shape().to_vec(),
470                view.as_slice()?.to_vec(),
471            ),
472            TensorView::C64(view) => <num_complex::Complex64 as TensorScalar>::into_tensor(
473                view.shape().to_vec(),
474                view.as_slice()?.to_vec(),
475            ),
476        }
477    }
478}
479
480/// Read-only retained gradient value.
481///
482/// # Examples
483///
484/// ```
485/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
486///
487/// let ctx = EagerRuntime::new()?;
488/// let x = EagerTensor::requires_grad_in(
489///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?,
490///     ctx,
491/// )?;
492/// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
493///     let squared = s.mul(&x, &x)?;
494///     s.reduce_sum(&squared, Some(&[0]))
495/// })?;
496/// let _gradients = loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
497/// let gradient = x.grad()?.expect("tracked leaf has a gradient");
498/// assert_eq!(gradient.shape(), &[2]);
499/// # Ok::<(), tenferro_ad::Error>(())
500/// ```
501#[derive(Clone, Debug)]
502pub struct GradientValue {
503    record: Arc<AdValueRecord>,
504}
505
506impl GradientValue {
507    /// Return the scalar dtype of the gradient.
508    pub fn dtype(&self) -> DType {
509        self.record.dtype()
510    }
511
512    /// Return the logical shape of the gradient.
513    pub fn shape(&self) -> &[usize] {
514        self.record.shape()
515    }
516
517    /// Borrow the gradient's value guard.
518    ///
519    /// # Errors
520    ///
521    /// Returns [`Error::RuntimeState`] when the retained gradient record is
522    /// unavailable or its allocation-group descriptor is invalid.
523    pub fn value(&self) -> Result<ValueGuard<'_>> {
524        self.record.value("GradientValue::value")
525    }
526
527    /// Borrow the gradient as a dtype-erased read target.
528    ///
529    /// # Errors
530    ///
531    /// Returns [`Error::RuntimeState`] when the retained gradient record or
532    /// its allocation-group descriptor is unavailable.
533    pub fn tensor_read(&self) -> Result<TensorRead<'_>> {
534        self.record.tensor_read("GradientValue::tensor_read")
535    }
536
537    /// Borrow a compact host slice without downloading backend storage.
538    ///
539    /// # Errors
540    ///
541    /// Returns [`Error::RuntimeState`] when the retained value is unavailable,
542    /// [`tenferro_tensor::ValidationError::DTypeMismatch`] when `T` does
543    /// not match the gradient dtype, or [`tenferro_tensor::Error::HostAccess`]
544    /// when backend storage cannot be mapped as a host slice.
545    pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&[T]> {
546        self.record
547            .value("GradientValue::as_slice")
548            .map_err(|error| {
549                tenferro_tensor::Error::runtime_state_source("GradientValue::as_slice", error)
550            })?
551            .as_slice()
552    }
553
554    /// Explicitly copy a host-resident gradient into a standalone tensor.
555    ///
556    /// # Errors
557    ///
558    /// Returns [`Error::RuntimeState`] when the retained value or execution
559    /// session is unavailable, or a typed backend/host-access error when the
560    /// gradient cannot be materialized as a contiguous tensor.
561    pub fn to_tensor(&self, engine: &mut dyn Engine) -> Result<Tensor> {
562        let value = self
563            .record
564            .value("GradientValue::to_tensor")
565            .map_err(|error| {
566                Error::runtime_state_source(
567                    "GradientValue::to_tensor",
568                    ErrorPhase::Execution,
569                    error,
570                )
571            })?;
572        match value.duplicate_host_tensor() {
573            Ok(tensor) => Ok(tensor),
574            Err(_) => {
575                let read = self.record.tensor_read("GradientValue::to_tensor")?;
576                engine.to_contiguous_read(read).map_err(Error::from)
577            }
578        }
579    }
580}
581
582/// Move-only accumulated gradient bundle backed by one allocation group.
583///
584/// # Examples
585///
586/// ```
587/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
588///
589/// let ctx = EagerRuntime::new()?;
590/// let x = EagerTensor::requires_grad_in(
591///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?,
592///     ctx,
593/// )?;
594/// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
595///     let squared = s.mul(&x, &x)?;
596///     s.reduce_sum(&squared, Some(&[0]))
597/// })?;
598/// let gradients = loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
599/// assert!(!gradients.is_empty());
600/// # Ok::<(), tenferro_ad::Error>(())
601/// ```
602#[derive(Debug)]
603pub struct Gradients {
604    group: AllocationGroup,
605    slots: HashMap<ValueKey<StdTensorOp>, DescriptorSlot>,
606}
607
608impl Gradients {
609    fn from_tensors(tensors: HashMap<ValueKey<StdTensorOp>, Tensor>) -> Result<Self> {
610        let (keys, values): (Vec<_>, Vec<_>) = tensors.into_iter().unzip();
611        let (group, bindings) = AllocationGroup::from_tensors(values).map_err(|error| {
612            Error::runtime_state_source("Gradients::from_tensors", ErrorPhase::Execution, error)
613        })?;
614        let slots = keys.into_iter().zip(bindings).collect();
615        Ok(Self { group, slots })
616    }
617
618    /// Return the number of retained gradient descriptors.
619    pub fn len(&self) -> usize {
620        self.slots.len()
621    }
622
623    /// Return whether no gradient was produced.
624    pub fn is_empty(&self) -> bool {
625        self.slots.is_empty()
626    }
627
628    /// Borrow one gradient view by its local value key.
629    pub fn grad(&self, key: &ValueKey<StdTensorOp>) -> Option<TensorView<'_>> {
630        let slot = self.slots.get(key).copied()?;
631        let mut reads = self.group.read_views(std::slice::from_ref(&slot)).ok()?;
632        match reads.pop()? {
633            TensorRead::View(view) => Some(view),
634            TensorRead::Tensor(_) => None,
635        }
636    }
637
638    /// Consume one gradient owner while leaving the bundle unchanged on failure.
639    ///
640    /// # Errors
641    ///
642    /// Returns [`tenferro_tensor::Error::RuntimeState`] when the descriptor is
643    /// invalid or its allocation is aliased. A missing key is reported as
644    /// `Ok(None)`.
645    pub fn take_grad(
646        &mut self,
647        key: &ValueKey<StdTensorOp>,
648    ) -> tenferro_tensor::Result<Option<Tensor>> {
649        let Some(&slot) = self.slots.get(key) else {
650            return Ok(None);
651        };
652        let tensor = self.group.take_tensor(slot).map_err(|error| {
653            tenferro_tensor::Error::runtime_state_source("Gradients::take_grad", error)
654        })?;
655        self.slots.remove(key);
656        Ok(Some(tensor))
657    }
658}
659
660/// Error returned when a value cannot be consumed without changing its owner.
661///
662/// # Examples
663///
664/// ```
665/// use tenferro_ad::{EagerRuntime, EagerTensor, IntoValueError, Tensor};
666///
667/// let ctx = EagerRuntime::new()?;
668/// let value = EagerTensor::from_tensor_in(
669///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?,
670///     ctx,
671/// )?;
672/// let _shared = value.clone();
673/// assert!(matches!(
674///     value.into_value(),
675///     Err(IntoValueError::NotUnique(_))
676/// ));
677/// # Ok::<(), tenferro_ad::Error>(())
678/// ```
679#[derive(Debug)]
680pub enum IntoValueError<H> {
681    /// Another eager handle, tape record, or checkpoint retains the value.
682    NotUnique(H),
683    /// Group extraction failed after the handle was uniquely acquired.
684    Extract { value: H, error: GroupError },
685}
686
687impl<H> std::fmt::Display for IntoValueError<H> {
688    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
689        match self {
690            Self::NotUnique(_) => formatter.write_str("eager value is retained by another handle"),
691            Self::Extract { error, .. } => {
692                write!(formatter, "eager value extraction failed: {error}")
693            }
694        }
695    }
696}
697
698impl<H: std::fmt::Debug + Send + Sync + 'static> std::error::Error for IntoValueError<H> {}
699
700/// What one direct retention container holds.
701// The pooled variant owns an allocation group inline; boxing it would add an
702// allocation to every retained value on the hot path.
703#[allow(clippy::large_enum_variant)]
704#[derive(Debug)]
705enum RetentionContainer {
706    /// Pooled storage owned through an allocation group and one descriptor slot.
707    Pooled {
708        group: AllocationGroup,
709        slot: DescriptorSlot,
710    },
711    /// A caller-owned value the runtime retains without taking pool ownership.
712    ///
713    /// The value returns to its owner when the record drops, and the runtime never
714    /// substitutes it for pooled storage. A caller-owned payload has no typed
715    /// descriptor view, so only its read and consume paths are available.
716    CallerOwned {
717        /// Boxed because a tensor value is much larger than the pooled variant.
718        tensor: Box<Tensor>,
719    },
720    /// An untracked result held directly.
721    ///
722    /// No AD group, residual or gradient will share it, so it needs no
723    /// allocation group: the tensor's own storage returns to its pool when the
724    /// record drops, and a unique handle hands the tensor back unchanged.
725    Owned { tensor: Tensor },
726}
727
728/// Read-only descriptor record used by eager handles and the AD registries.
729#[derive(Debug)]
730pub(crate) struct AdValueRecord {
731    container: Arc<RetentionContainer>,
732    dtype: DType,
733    shape: Box<[usize]>,
734}
735
736impl AdValueRecord {
737    fn from_group(
738        group: AllocationGroup,
739        slot: DescriptorSlot,
740        dtype: DType,
741        shape: Vec<usize>,
742    ) -> Arc<Self> {
743        Arc::new(Self {
744            container: Arc::new(RetentionContainer::Pooled { group, slot }),
745            dtype,
746            shape: shape.into_boxed_slice(),
747        })
748    }
749
750    fn from_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
751        let dtype = tensor.dtype();
752        let shape = tensor.shape().to_vec();
753        if matches!(dtype, DType::External(_)) {
754            // A caller-owned payload is retained directly: it owns no pooled
755            // storage, so there is no group to build and nothing to return to a
756            // pool when the record drops.
757            return Ok(Arc::new(Self {
758                container: Arc::new(RetentionContainer::CallerOwned {
759                    tensor: Box::new(tensor),
760                }),
761                dtype,
762                shape: shape.into_boxed_slice(),
763            }));
764        }
765        let (group, bindings) = AllocationGroup::from_tensors(vec![tensor])
766            .map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
767        let slot = bindings.first().copied().ok_or_else(|| {
768            Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
769        })?;
770        Ok(Self::from_group(group, slot, dtype, shape))
771    }
772
773    /// Retain an untracked result without building an allocation group.
774    fn from_untracked_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
775        if matches!(tensor.dtype(), DType::External(_)) {
776            return Self::from_tensor(tensor, op);
777        }
778        let dtype = tensor.dtype();
779        let shape = tensor.shape().to_vec().into_boxed_slice();
780        Ok(Arc::new(Self {
781            container: Arc::new(RetentionContainer::Owned { tensor }),
782            dtype,
783            shape,
784        }))
785    }
786
787    fn tensor_read(&self, op: &'static str) -> Result<TensorRead<'_>> {
788        match self.container.as_ref() {
789            RetentionContainer::Pooled { group, slot } => group.read_view(*slot).map_err(|error| {
790                // One slot, so no `Vec` and no empty-result case: the single-slot
791                // accessor is the whole request (#2079).
792                Error::runtime_state_source(op, ErrorPhase::Execution, error)
793            }),
794            RetentionContainer::CallerOwned { tensor } => Ok(TensorRead::from_tensor(tensor)),
795            RetentionContainer::Owned { tensor } => Ok(TensorRead::from_tensor(tensor)),
796        }
797    }
798
799    fn value(&self, op: &'static str) -> Result<ValueGuard<'_>> {
800        if let RetentionContainer::Owned { tensor } = self.container.as_ref() {
801            // A preset-dtype owned tensor always has a typed borrowed view.
802            return Ok(ValueGuard {
803                view: TensorRead::from_tensor(tensor).tensor_view(),
804            });
805        }
806        match self.tensor_read(op)? {
807            TensorRead::View(view) => Ok(ValueGuard { view }),
808            // A caller-owned payload has no typed descriptor view, so a path that
809            // needs one fails explicitly instead of borrowing the payload as bytes.
810            TensorRead::Tensor(_) => Err(Error::runtime_state(
811                op,
812                ErrorPhase::Execution,
813                "allocation-group value did not produce a borrowed descriptor view",
814            )),
815        }
816    }
817
818    fn dtype(&self) -> DType {
819        self.dtype
820    }
821
822    fn shape(&self) -> &[usize] {
823        &self.shape
824    }
825}
826
827/// Shared eager execution context for tensors on a backend.
828///
829/// Reusing one context lets eager tensors share backend state, extension
830/// runtime caches, and gradient storage across a computation.
831///
832/// # Examples
833///
834/// ```
835/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
836///
837/// let ctx = EagerRuntime::new()?;
838/// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
839/// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx.clone()).unwrap();
840/// let z = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.add(&x, &y)).unwrap();
841///
842/// assert_eq!(z.value().unwrap().as_slice::<f64>().unwrap(), &[3.0]);
843/// # Ok::<(), tenferro_ad::Error>(())
844/// ```
845pub struct EagerRuntime {
846    id: ContextId,
847    /// Prepared-program runtime. Engine kinds are registered on first use;
848    /// the runtime never stores an engine.
849    runtime: Runtime,
850    /// Serializes runtime reconfiguration (engine-kind registration and
851    /// extension-module installation).
852    extension_install_lock: Mutex<()>,
853    semantic_extension_rules: SemanticExtensionRuleSet,
854    grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
855    value_records: Mutex<HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>,
856    ad_transform_cache: Arc<AdTransformCache>,
857    /// S2: prepared derivative programs keyed by semantic structure, wrt input,
858    /// and concrete bound input metadata. Avoids re-running freeze+AD
859    /// transform+compile_frozen on warm structure hits.
860    prepared_derivative_cache: Mutex<PreparedDerivativeCache>,
861}
862
863/// An eager runtime and its borrowed backend session for one execution boundary.
864///
865/// Obtain this only through [`EagerRuntime::with_eager_session`]. It rejects
866/// tensors from another eager runtime even if both runtimes use the same backend
867/// type, and it cannot escape the boundary closure.
868///
869/// # Examples
870///
871/// ```rust
872/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
873///
874/// let ctx = EagerRuntime::new()?;
875/// let x = EagerTensor::from_tensor_in(
876///     Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?, ctx.clone(),
877/// )?;
878/// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.neg(&x))?;
879/// assert_eq!(y.value()?.as_slice::<f64>()?, &[-3.0]);
880/// # Ok::<(), tenferro_ad::Error>(())
881/// ```
882///
883/// The borrowed session cannot escape its execution boundary:
884///
885/// ```compile_fail
886/// use tenferro_ad::EagerRuntime;
887/// let ctx = EagerRuntime::new().unwrap();
888/// let escaped = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session).unwrap();
889/// let _ = escaped;
890/// ```
891pub struct EagerSession<'a> {
892    runtime: &'a Arc<EagerRuntime>,
893    engine: &'a mut dyn ExecutionEngine,
894}
895
896impl fmt::Debug for EagerSession<'_> {
897    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
898        f.debug_struct("EagerSession")
899            .field("runtime_id", &self.runtime.id())
900            .finish_non_exhaustive()
901    }
902}
903
904impl EagerSession<'_> {
905    /// Negate an eager tensor inside the caller's execution boundary.
906    ///
907    /// # Examples
908    ///
909    /// ```rust
910    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
911    ///
912    /// let ctx = EagerRuntime::new()?;
913    /// let x = EagerTensor::from_tensor_in(
914    ///     Tensor::from_vec_col_major(vec![1], vec![4.0_f64])?, ctx.clone(),
915    /// )?;
916    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.neg(&x))?;
917    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-4.0]);
918    /// # Ok::<(), tenferro_ad::Error>(())
919    /// ```
920    ///
921    /// # Errors
922    ///
923    /// Returns [`Error::ContextMismatch`] for a tensor from another runtime,
924    /// or a typed eager/backend error from the selected operation.
925    pub fn neg(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
926        self.run_unary(input, StdTensorOp::Neg)
927    }
928
929    /// Compute the elementwise exponential inside this eager session.
930    ///
931    /// # Examples
932    ///
933    /// ```rust
934    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
935    /// let ctx = EagerRuntime::new()?;
936    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?, ctx.clone())?;
937    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.exp(&x))?;
938    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0]);
939    /// # Ok::<(), tenferro_ad::Error>(())
940    /// ```
941    ///
942    /// # Errors
943    ///
944    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
945    /// unsupported/backend error for the input dtype.
946    pub fn exp(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
947        self.run_unary(input, StdTensorOp::Exp)
948    }
949
950    /// Compute the elementwise absolute value inside this eager session.
951    ///
952    /// # Examples
953    ///
954    /// ```rust
955    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
956    /// let ctx = EagerRuntime::new()?;
957    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![-2.0_f64])?, ctx.clone())?;
958    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.abs(&x))?;
959    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
960    /// # Ok::<(), tenferro_ad::Error>(())
961    /// ```
962    ///
963    /// # Errors
964    ///
965    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
966    /// unsupported/backend error for the input dtype.
967    pub fn abs(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
968        self.run_unary(input, StdTensorOp::Abs)
969    }
970
971    /// Compute the elementwise conjugate inside this eager session.
972    ///
973    /// # Examples
974    ///
975    /// ```rust
976    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
977    /// let ctx = EagerRuntime::new()?;
978    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?, ctx.clone())?;
979    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.conj(&x))?;
980    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
981    /// # Ok::<(), tenferro_ad::Error>(())
982    /// ```
983    ///
984    /// # Errors
985    ///
986    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
987    /// unsupported/backend error for the input dtype.
988    pub fn conj(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
989        self.run_unary(input, StdTensorOp::Conj)
990    }
991
992    /// Compute the elementwise sign on this borrowed session.
993    ///
994    /// # Examples
995    /// ```rust
996    /// use tenferro_ad::{EagerRuntime, Tensor};
997    /// let ctx = EagerRuntime::new()?;
998    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
999    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![-2.0_f64])?)?;
1000    ///     s.sign(&x)
1001    /// })?;
1002    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-1.0]);
1003    /// # Ok::<(), tenferro_ad::Error>(())
1004    /// ```
1005    /// # Errors
1006    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1007    pub fn sign(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1008        self.run_unary(input, StdTensorOp::Sign)
1009    }
1010
1011    /// Compute the elementwise natural logarithm on this borrowed session.
1012    ///
1013    /// # Examples
1014    /// ```rust
1015    /// use tenferro_ad::{EagerRuntime, Tensor};
1016    /// let ctx = EagerRuntime::new()?;
1017    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1018    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?)?;
1019    ///     s.log(&x)
1020    /// })?;
1021    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1022    /// # Ok::<(), tenferro_ad::Error>(())
1023    /// ```
1024    /// # Errors
1025    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1026    pub fn log(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1027        self.run_unary(input, StdTensorOp::Log)
1028    }
1029
1030    /// Compute the elementwise square root on this borrowed session.
1031    ///
1032    /// # Examples
1033    /// ```rust
1034    /// use tenferro_ad::{EagerRuntime, Tensor};
1035    /// let ctx = EagerRuntime::new()?;
1036    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1037    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![4.0_f64])?)?;
1038    ///     s.sqrt(&x)
1039    /// })?;
1040    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
1041    /// # Ok::<(), tenferro_ad::Error>(())
1042    /// ```
1043    /// # Errors
1044    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1045    pub fn sqrt(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1046        self.run_unary(input, StdTensorOp::Sqrt)
1047    }
1048
1049    /// Compute the elementwise reciprocal square root on this borrowed session.
1050    ///
1051    /// # Examples
1052    /// ```rust
1053    /// use tenferro_ad::{EagerRuntime, Tensor};
1054    /// let ctx = EagerRuntime::new()?;
1055    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1056    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![4.0_f64])?)?;
1057    ///     s.rsqrt(&x)
1058    /// })?;
1059    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.5]);
1060    /// # Ok::<(), tenferro_ad::Error>(())
1061    /// ```
1062    /// # Errors
1063    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1064    pub fn rsqrt(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1065        self.run_unary(input, StdTensorOp::Rsqrt)
1066    }
1067
1068    /// Compute the elementwise sine on this borrowed session.
1069    ///
1070    /// # Examples
1071    /// ```rust
1072    /// use tenferro_ad::{EagerRuntime, Tensor};
1073    /// let ctx = EagerRuntime::new()?;
1074    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1075    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1076    ///     s.sin(&x)
1077    /// })?;
1078    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1079    /// # Ok::<(), tenferro_ad::Error>(())
1080    /// ```
1081    /// # Errors
1082    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1083    pub fn sin(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1084        self.run_unary(input, StdTensorOp::Sin)
1085    }
1086
1087    /// Compute the elementwise cosine on this borrowed session.
1088    ///
1089    /// # Examples
1090    /// ```rust
1091    /// use tenferro_ad::{EagerRuntime, Tensor};
1092    /// let ctx = EagerRuntime::new()?;
1093    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1094    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1095    ///     s.cos(&x)
1096    /// })?;
1097    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0]);
1098    /// # Ok::<(), tenferro_ad::Error>(())
1099    /// ```
1100    /// # Errors
1101    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1102    pub fn cos(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1103        self.run_unary(input, StdTensorOp::Cos)
1104    }
1105
1106    /// Compute the elementwise hyperbolic tangent on this borrowed session.
1107    ///
1108    /// # Examples
1109    /// ```rust
1110    /// use tenferro_ad::{EagerRuntime, Tensor};
1111    /// let ctx = EagerRuntime::new()?;
1112    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1113    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1114    ///     s.tanh(&x)
1115    /// })?;
1116    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1117    /// # Ok::<(), tenferro_ad::Error>(())
1118    /// ```
1119    /// # Errors
1120    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1121    pub fn tanh(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1122        self.run_unary(input, StdTensorOp::Tanh)
1123    }
1124
1125    /// Compute `exp(x) - 1` elementwise on this borrowed session.
1126    ///
1127    /// # Examples
1128    /// ```rust
1129    /// use tenferro_ad::{EagerRuntime, Tensor};
1130    /// let ctx = EagerRuntime::new()?;
1131    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1132    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1133    ///     s.expm1(&x)
1134    /// })?;
1135    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1136    /// # Ok::<(), tenferro_ad::Error>(())
1137    /// ```
1138    /// # Errors
1139    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1140    pub fn expm1(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1141        self.run_unary(input, StdTensorOp::Expm1)
1142    }
1143
1144    /// Compute `log(1 + x)` elementwise on this borrowed session.
1145    ///
1146    /// # Examples
1147    /// ```rust
1148    /// use tenferro_ad::{EagerRuntime, Tensor};
1149    /// let ctx = EagerRuntime::new()?;
1150    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1151    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1152    ///     s.log1p(&x)
1153    /// })?;
1154    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1155    /// # Ok::<(), tenferro_ad::Error>(())
1156    /// ```
1157    /// # Errors
1158    /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1159    pub fn log1p(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1160        self.run_unary(input, StdTensorOp::Log1p)
1161    }
1162
1163    /// Compute the error function `erf(x)` elementwise on this borrowed session.
1164    ///
1165    /// Defined for real `F32`/`F64` tensors; `erf(+-0) = +-0`,
1166    /// `erf(+-inf) = +-1`, and `NaN` stays `NaN`. The derivative is
1167    /// `2/sqrt(pi) * exp(-x^2)`.
1168    ///
1169    /// # Examples
1170    /// ```rust
1171    /// use tenferro_ad::{EagerRuntime, Tensor};
1172    /// let ctx = EagerRuntime::new()?;
1173    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1174    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![0.0_f64, 1.0])?)?;
1175    ///     s.erf(&x)
1176    /// })?;
1177    /// let y = y.value()?;
1178    /// let y = y.as_slice::<f64>()?;
1179    /// assert_eq!(y[0], 0.0);
1180    /// assert!((y[1] - 0.842_700_792_949_714_9).abs() < 1.0e-15);
1181    /// # Ok::<(), tenferro_ad::Error>(())
1182    /// ```
1183    /// # Errors
1184    /// Returns a typed foreign-runtime error, a typed unsupported-dtype error
1185    /// for complex, integer, or `Bool` input, or a backend error.
1186    pub fn erf(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1187        self.run_unary(input, StdTensorOp::Erf)
1188    }
1189
1190    /// Convert a tensor under the checked dtype-promotion lattice.
1191    /// Use [`Self::cast`] for intentional lossy projection.
1192    ///
1193    /// # Examples
1194    ///
1195    /// ```rust
1196    /// use tenferro_ad::{DType, EagerRuntime, Tensor};
1197    /// let ctx = EagerRuntime::new()?;
1198    /// let converted = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1199    ///     let x = session.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
1200    ///     session.convert(&x, DType::C64)
1201    /// })?;
1202    /// assert_eq!(converted.dtype(), DType::C64);
1203    /// # Ok::<(), tenferro_ad::Error>(())
1204    /// ```
1205    ///
1206    /// # Errors
1207    ///
1208    /// Returns [`Error::ContextMismatch`] for a foreign runtime or a typed
1209    /// unsupported dtype conversion/backend error.
1210    pub fn convert(&mut self, input: &EagerTensor, to: DType) -> Result<EagerTensor> {
1211        self.ensure_runtime(input)?;
1212        tenferro_tensor::validate::validate_convert_dtype(
1213            "EagerTensor::convert",
1214            input.dtype(),
1215            to,
1216        )
1217        .map_err(Error::TensorRuntime)?;
1218        self.cast(input, to)
1219    }
1220
1221    /// Cast a tensor to a dtype, permitting explicitly lossy projections.
1222    ///
1223    /// # Examples
1224    ///
1225    /// ```rust
1226    /// use tenferro_ad::{DType, EagerRuntime, Tensor};
1227    /// let ctx = EagerRuntime::new()?;
1228    /// let casted = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1229    ///     let x = session.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.8_f64])?)?;
1230    ///     session.cast(&x, DType::I32)
1231    /// })?;
1232    /// assert_eq!(casted.value()?.as_slice::<i32>()?, &[2]);
1233    /// # Ok::<(), tenferro_ad::Error>(())
1234    /// ```
1235    ///
1236    /// # Errors
1237    ///
1238    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1239    /// unsupported projection/backend error.
1240    pub fn cast(&mut self, input: &EagerTensor, to: DType) -> Result<EagerTensor> {
1241        self.run_unary(
1242            input,
1243            StdTensorOp::Convert {
1244                from: input.dtype(),
1245                to,
1246            },
1247        )
1248    }
1249
1250    /// Permute the axes of an eager tensor while preserving independent ownership.
1251    ///
1252    /// # Examples
1253    ///
1254    /// ```rust
1255    /// use tenferro_ad::{EagerRuntime, Tensor};
1256    /// let ctx = EagerRuntime::new()?;
1257    /// let copied = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1258    ///     let x = session.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1259    ///     let y = session.transpose(&x, &[1, 0])?;
1260    ///     session.duplicate_value(&y)
1261    /// })?;
1262    /// assert_eq!(copied.as_slice::<f64>()?, &[1.0, 3.0, 2.0, 4.0]);
1263    /// # Ok::<(), tenferro_ad::Error>(())
1264    /// ```
1265    ///
1266    /// # Errors
1267    ///
1268    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1269    /// axis/backend error for an invalid permutation or copy.
1270    pub fn transpose(&mut self, input: &EagerTensor, perm: &[usize]) -> Result<EagerTensor> {
1271        self.ensure_runtime(input)?;
1272        input.transpose_in_session(perm, &mut *self.engine)
1273    }
1274
1275    /// Reshape an eager tensor while retaining a separate result owner.
1276    ///
1277    /// # Examples
1278    ///
1279    /// ```rust
1280    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1281    /// let ctx = EagerRuntime::new()?;
1282    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1283    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.reshape(&x, [1, 2]))?;
1284    /// assert_eq!(y.shape(), &[1, 2]);
1285    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0, 2.0]);
1286    /// # Ok::<(), tenferro_ad::Error>(())
1287    /// ```
1288    ///
1289    /// # Errors
1290    ///
1291    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1292    /// validation/backend error when the target shape is incompatible.
1293    pub fn reshape(
1294        &mut self,
1295        input: &EagerTensor,
1296        shape: impl IntoShapeVec,
1297    ) -> Result<EagerTensor> {
1298        self.ensure_runtime(input)?;
1299        input.reshape_in_session(&shape.into_shape_vec(), &mut *self.engine)
1300    }
1301
1302    /// Slice an eager tensor with explicit start, limit, and stride per axis.
1303    ///
1304    /// # Examples
1305    ///
1306    /// ```rust
1307    /// use tenferro_ad::{EagerRuntime, SliceConfig, Tensor};
1308    /// let ctx = EagerRuntime::new()?;
1309    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1310    ///     let x = session.constant_from(Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1311    ///     session.slice(&x, SliceConfig { starts: vec![1], limits: vec![3], strides: vec![1] })
1312    /// })?;
1313    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0, 3.0]);
1314    /// # Ok::<(), tenferro_ad::Error>(())
1315    /// ```
1316    ///
1317    /// # Errors
1318    ///
1319    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1320    /// axis/stride/backend error for an invalid slice or copy.
1321    pub fn slice(&mut self, input: &EagerTensor, config: SliceConfig) -> Result<EagerTensor> {
1322        self.ensure_runtime(input)?;
1323        input.slice_in_session(config, &mut *self.engine)
1324    }
1325
1326    /// Broadcast an eager tensor into a larger shape on this session.
1327    ///
1328    /// # Examples
1329    ///
1330    /// ```rust
1331    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1332    /// let ctx = EagerRuntime::new()?;
1333    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1334    /// let copy = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1335    ///     let y = session.broadcast_in_dim(&x, &[2, 2], &[0])?;
1336    ///     session.duplicate_value(&y)
1337    /// })?;
1338    /// assert_eq!(copy.as_slice::<f64>()?, &[1.0, 2.0, 1.0, 2.0]);
1339    /// # Ok::<(), tenferro_ad::Error>(())
1340    /// ```
1341    ///
1342    /// # Errors
1343    ///
1344    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1345    /// validation/backend error for an invalid broadcast mapping.
1346    pub fn broadcast_in_dim(
1347        &mut self,
1348        input: &EagerTensor,
1349        shape: &[usize],
1350        dims: &[usize],
1351    ) -> Result<EagerTensor> {
1352        self.ensure_runtime(input)?;
1353        input.broadcast_in_dim_in_session(shape, dims, &mut *self.engine)
1354    }
1355
1356    /// Keep the lower triangle of a matrix on this borrowed session.
1357    ///
1358    /// # Examples
1359    /// ```rust
1360    /// use tenferro_ad::{EagerRuntime, Tensor};
1361    /// let ctx = EagerRuntime::new()?;
1362    /// let lower = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1363    ///     let matrix = s.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1364    ///     s.tril(&matrix, 0)
1365    /// })?;
1366    /// assert_eq!(lower.value()?.as_slice::<f64>()?, &[1.0, 2.0, 0.0, 4.0]);
1367    /// # Ok::<(), tenferro_ad::Error>(())
1368    /// ```
1369    /// # Errors
1370    /// Returns a typed foreign-runtime, rank, unsupported-dtype, or backend error.
1371    pub fn tril(&mut self, input: &EagerTensor, k: i64) -> Result<EagerTensor> {
1372        self.run_unary(input, StdTensorOp::Tril { k })
1373    }
1374
1375    /// Keep the upper triangle of a matrix on this borrowed session.
1376    ///
1377    /// # Examples
1378    /// ```rust
1379    /// use tenferro_ad::{EagerRuntime, Tensor};
1380    /// let ctx = EagerRuntime::new()?;
1381    /// let upper = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1382    ///     let matrix = s.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1383    ///     s.triu(&matrix, 0)
1384    /// })?;
1385    /// assert_eq!(upper.value()?.as_slice::<f64>()?, &[1.0, 0.0, 3.0, 4.0]);
1386    /// # Ok::<(), tenferro_ad::Error>(())
1387    /// ```
1388    /// # Errors
1389    /// Returns a typed foreign-runtime, rank, unsupported-dtype, or backend error.
1390    pub fn triu(&mut self, input: &EagerTensor, k: i64) -> Result<EagerTensor> {
1391        self.run_unary(input, StdTensorOp::Triu { k })
1392    }
1393
1394    /// Pad an eager tensor with zeros on this borrowed session.
1395    ///
1396    /// # Examples
1397    /// ```rust
1398    /// use tenferro_ad::{EagerRuntime, PadConfig, Tensor};
1399    /// let ctx = EagerRuntime::new()?;
1400    /// let padded = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1401    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
1402    ///     s.pad(&x, PadConfig {
1403    ///         edge_padding_low: vec![1],
1404    ///         edge_padding_high: vec![1],
1405    ///         interior_padding: vec![1],
1406    ///     })
1407    /// })?;
1408    /// assert_eq!(padded.value()?.as_slice::<f64>()?, &[0.0, 1.0, 0.0, 2.0, 0.0]);
1409    /// # Ok::<(), tenferro_ad::Error>(())
1410    /// ```
1411    /// # Errors
1412    /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
1413    /// runtime, a validation error with
1414    /// `ValidationError::InvalidArgument` for a padding configuration whose
1415    /// length or extents do not match the input rank, or
1416    /// [`Error::TensorRuntime`] for a typed backend failure.
1417    pub fn pad(&mut self, input: &EagerTensor, config: PadConfig) -> Result<EagerTensor> {
1418        self.run_unary(input, StdTensorOp::Pad(config))
1419    }
1420
1421    /// Reverse the elements along selected axes on this borrowed session.
1422    ///
1423    /// # Examples
1424    /// ```rust
1425    /// use tenferro_ad::{EagerRuntime, Tensor};
1426    /// let ctx = EagerRuntime::new()?;
1427    /// let reversed = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1428    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0])?)?;
1429    ///     s.reverse(&x, &[0])
1430    /// })?;
1431    /// assert_eq!(reversed.value()?.as_slice::<f64>()?, &[3.0, 2.0, 1.0]);
1432    /// # Ok::<(), tenferro_ad::Error>(())
1433    /// ```
1434    /// # Errors
1435    /// Returns a typed foreign-runtime, invalid-axis, or backend error.
1436    pub fn reverse(&mut self, input: &EagerTensor, axes: &[usize]) -> Result<EagerTensor> {
1437        self.ensure_runtime(input)?;
1438        crate::eager_ops::validate_eager_axes("EagerSession::reverse", input.shape().len(), axes)?;
1439        self.run_unary(
1440            input,
1441            StdTensorOp::Reverse {
1442                axes: axes.to_vec(),
1443            },
1444        )
1445    }
1446
1447    /// Slice an eager tensor using runtime start indices in this borrowed session.
1448    ///
1449    /// # Examples
1450    /// ```rust
1451    /// use tenferro_ad::{EagerRuntime, Tensor};
1452    /// let ctx = EagerRuntime::new()?;
1453    /// let selected = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1454    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1455    ///     let starts = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![1_i64])?)?;
1456    ///     s.dynamic_slice(&x, &starts, &[2])
1457    /// })?;
1458    /// assert_eq!(selected.value()?.as_slice::<f64>()?, &[2.0, 3.0]);
1459    /// # Ok::<(), tenferro_ad::Error>(())
1460    /// ```
1461    /// # Errors
1462    /// Returns a typed foreign-runtime, invalid-index or slice-shape, or backend error.
1463    pub fn dynamic_slice(
1464        &mut self,
1465        input: &EagerTensor,
1466        starts: &EagerTensor,
1467        sizes: &[usize],
1468    ) -> Result<EagerTensor> {
1469        self.ensure_runtime(input)?;
1470        self.ensure_runtime(starts)?;
1471        EagerTensor::nary_op_in_session(
1472            &[input, starts],
1473            StdTensorOp::DynamicSlice {
1474                slice_sizes: sizes.to_vec(),
1475            },
1476            &mut *self.engine,
1477        )
1478    }
1479
1480    /// Gather elements of an eager tensor in this borrowed session.
1481    ///
1482    /// # Examples
1483    /// ```rust
1484    /// use tenferro_ad::{EagerRuntime, GatherConfig, Tensor};
1485    /// let ctx = EagerRuntime::new()?;
1486    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1487    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0])?)?;
1488    ///     let indices = s.constant_from(Tensor::from_vec_col_major(vec![2, 1], vec![2_i64, 0])?)?;
1489    ///     s.gather(&x, &indices, GatherConfig {
1490    ///         offset_dims: vec![], collapsed_slice_dims: vec![0],
1491    ///         start_index_map: vec![0], index_vector_dim: 1,
1492    ///         slice_sizes: vec![1],
1493    ///     })
1494    /// })?;
1495    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[30.0, 10.0]);
1496    /// # Ok::<(), tenferro_ad::Error>(())
1497    /// ```
1498    /// # Errors
1499    /// Returns a typed foreign-runtime, invalid-index/configuration, or backend error.
1500    pub fn gather(
1501        &mut self,
1502        input: &EagerTensor,
1503        indices: &EagerTensor,
1504        config: GatherConfig,
1505    ) -> Result<EagerTensor> {
1506        self.ensure_runtime(input)?;
1507        self.ensure_runtime(indices)?;
1508        EagerTensor::nary_op_in_session(
1509            &[input, indices],
1510            StdTensorOp::Gather(config),
1511            &mut *self.engine,
1512        )
1513    }
1514
1515    /// Concatenate eager tensors along one axis in this borrowed session.
1516    ///
1517    /// # Examples
1518    /// ```rust
1519    /// use tenferro_ad::{EagerRuntime, Tensor};
1520    /// let ctx = EagerRuntime::new()?;
1521    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1522    ///     let a = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?)?;
1523    ///     let b = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
1524    ///     s.concatenate(&[&a, &b], 0)
1525    /// })?;
1526    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[1.0, 2.0]);
1527    /// # Ok::<(), tenferro_ad::Error>(())
1528    /// ```
1529    /// # Errors
1530    /// Returns a typed empty-input, foreign-runtime, invalid-axis/shape, or backend error.
1531    pub fn concatenate(&mut self, inputs: &[&EagerTensor], axis: usize) -> Result<EagerTensor> {
1532        for input in inputs {
1533            self.ensure_runtime(input)?;
1534        }
1535        EagerTensor::nary_op_in_session(
1536            inputs,
1537            StdTensorOp::Concatenate {
1538                axis,
1539                input_count: inputs.len(),
1540            },
1541            &mut *self.engine,
1542        )
1543    }
1544
1545    /// Scatter updates into an eager tensor within this borrowed session.
1546    ///
1547    /// # Examples
1548    /// ```rust
1549    /// use tenferro_ad::{EagerRuntime, ScatterConfig, Tensor};
1550    /// let ctx = EagerRuntime::new()?;
1551    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1552    ///     let input = s.constant_from(Tensor::from_vec_col_major(vec![4], vec![0.0_f64; 4])?)?;
1553    ///     let indices = s.constant_from(Tensor::from_vec_col_major(vec![2, 1], vec![1_i64, 3])?)?;
1554    ///     let updates = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![5.0_f64, 7.0])?)?;
1555    ///     s.scatter(&input, &indices, &updates, ScatterConfig {
1556    ///         update_window_dims: vec![],
1557    ///         inserted_window_dims: vec![0],
1558    ///         scatter_dims_to_operand_dims: vec![0],
1559    ///         index_vector_dim: 1,
1560    ///     })
1561    /// })?;
1562    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[0.0, 5.0, 0.0, 7.0]);
1563    /// # Ok::<(), tenferro_ad::Error>(())
1564    /// ```
1565    /// # Errors
1566    /// Returns a typed foreign-runtime, invalid-index/configuration, or backend error.
1567    pub fn scatter(
1568        &mut self,
1569        input: &EagerTensor,
1570        indices: &EagerTensor,
1571        updates: &EagerTensor,
1572        config: ScatterConfig,
1573    ) -> Result<EagerTensor> {
1574        self.ensure_runtime(input)?;
1575        self.ensure_runtime(indices)?;
1576        self.ensure_runtime(updates)?;
1577        EagerTensor::nary_op_in_session(
1578            &[input, indices, updates],
1579            StdTensorOp::Scatter(config),
1580            &mut *self.engine,
1581        )
1582    }
1583
1584    /// Extract a diagonal along two axes in this borrowed session.
1585    ///
1586    /// # Examples
1587    /// ```rust
1588    /// use tenferro_ad::{EagerRuntime, Tensor};
1589    /// let ctx = EagerRuntime::new()?;
1590    /// let diagonal = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1591    ///     let matrix = s.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1592    ///     s.extract_diag(&matrix, 0, 1)
1593    /// })?;
1594    /// assert_eq!(diagonal.value()?.as_slice::<f64>()?, &[1.0, 4.0]);
1595    /// # Ok::<(), tenferro_ad::Error>(())
1596    /// ```
1597    /// # Errors
1598    /// Returns a typed foreign-runtime, invalid-axis, or backend error.
1599    pub fn extract_diag(
1600        &mut self,
1601        input: &EagerTensor,
1602        axis_a: usize,
1603        axis_b: usize,
1604    ) -> Result<EagerTensor> {
1605        self.ensure_runtime(input)?;
1606        self.run_unary(input, StdTensorOp::ExtractDiag { axis_a, axis_b })
1607    }
1608
1609    /// Embed the input along a diagonal in this borrowed session.
1610    ///
1611    /// # Examples
1612    /// ```rust
1613    /// use tenferro_ad::{EagerRuntime, Tensor};
1614    /// let ctx = EagerRuntime::new()?;
1615    /// let matrix = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1616    ///     let diagonal = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
1617    ///     s.embed_diag(&diagonal, 0, 1)
1618    /// })?;
1619    /// assert_eq!(matrix.value()?.as_slice::<f64>()?, &[1.0, 0.0, 0.0, 2.0]);
1620    /// # Ok::<(), tenferro_ad::Error>(())
1621    /// ```
1622    /// # Errors
1623    /// Returns a typed foreign-runtime, invalid-axis, or backend error.
1624    pub fn embed_diag(
1625        &mut self,
1626        input: &EagerTensor,
1627        axis_a: usize,
1628        axis_b: usize,
1629    ) -> Result<EagerTensor> {
1630        self.ensure_runtime(input)?;
1631        self.run_unary(input, StdTensorOp::EmbedDiag { axis_a, axis_b })
1632    }
1633
1634    /// Reduce an eager tensor over selected axes within this borrowed session.
1635    /// `None` reduces all axes, while `Some(&[])` retains the input shape.
1636    ///
1637    /// # Examples
1638    ///
1639    /// ```rust
1640    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1641    /// let ctx = EagerRuntime::new()?;
1642    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1643    /// let sum = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.reduce_sum(&x, None))?;
1644    /// assert_eq!(sum.value()?.as_slice::<f64>()?, &[3.0]);
1645    /// # Ok::<(), tenferro_ad::Error>(())
1646    /// ```
1647    ///
1648    /// # Errors
1649    ///
1650    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1651    /// validation/backend error for invalid axes or unsupported dtypes.
1652    pub fn reduce_sum(
1653        &mut self,
1654        input: &EagerTensor,
1655        axes: Option<&[usize]>,
1656    ) -> Result<EagerTensor> {
1657        self.ensure_runtime(input)?;
1658        input.reduce_sum_in_session(axes, &mut *self.engine)
1659    }
1660
1661    /// Sum elementwise squares over the selected axes in this borrowed session.
1662    /// Only `f32` and `f64` are supported. `None` reduces every axis, like the
1663    /// rest of the reduction family; `Some(&[])` squares each value.
1664    ///
1665    /// # Examples
1666    /// ```rust
1667    /// use tenferro_ad::{EagerRuntime, Tensor};
1668    /// let ctx = EagerRuntime::new()?;
1669    /// let (sum, all) = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1670    ///     let input = s.constant_from(Tensor::from_vec_col_major([2], vec![3.0_f64, 4.0])?)?;
1671    ///     Ok::<_, tenferro_ad::Error>((
1672    ///         s.reduce_sum_squares(&input, Some(&[0]))?,
1673    ///         s.reduce_sum_squares(&input, None)?,
1674    ///     ))
1675    /// })?;
1676    /// assert_eq!(sum.value()?.as_slice::<f64>()?, &[25.0]);
1677    /// assert_eq!(all.value()?.as_slice::<f64>()?, &[25.0]);
1678    /// # Ok::<(), tenferro_ad::Error>(())
1679    /// ```
1680    /// # Errors
1681    /// Returns typed foreign-runtime, invalid-axis, unsupported-dtype, or backend errors.
1682    pub fn reduce_sum_squares(
1683        &mut self,
1684        input: &EagerTensor,
1685        axes: Option<&[usize]>,
1686    ) -> Result<EagerTensor> {
1687        self.ensure_runtime(input)?;
1688        let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1689        crate::eager_ops::validate_eager_axes(
1690            "EagerSession::reduce_sum_squares",
1691            input.shape().len(),
1692            &axes,
1693        )?;
1694        self.run_unary(input, StdTensorOp::ReduceSumSquares { axes })
1695    }
1696
1697    /// Reduce the product of selected axes in this borrowed session.
1698    /// `None` reduces every axis.
1699    ///
1700    /// # Examples
1701    /// ```rust
1702    /// use tenferro_ad::{EagerRuntime, Tensor};
1703    /// let ctx = EagerRuntime::new()?;
1704    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1705    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?)?;
1706    ///     s.reduce_prod(&x, None)
1707    /// })?;
1708    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[6.0]);
1709    /// # Ok::<(), tenferro_ad::Error>(())
1710    /// ```
1711    /// # Errors
1712    /// Returns a typed foreign-runtime, invalid-axis, unsupported-dtype, or backend error.
1713    pub fn reduce_prod(
1714        &mut self,
1715        input: &EagerTensor,
1716        axes: Option<&[usize]>,
1717    ) -> Result<EagerTensor> {
1718        self.ensure_runtime(input)?;
1719        let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1720        crate::eager_ops::validate_eager_axes(
1721            "EagerSession::reduce_prod",
1722            input.shape().len(),
1723            &axes,
1724        )?;
1725        self.run_unary(input, StdTensorOp::ReduceProd { axes })
1726    }
1727
1728    /// Reduce the maximum over selected axes in this borrowed session.
1729    /// `None` reduces every axis.
1730    ///
1731    /// # Examples
1732    /// ```rust
1733    /// use tenferro_ad::{EagerRuntime, Tensor};
1734    /// let ctx = EagerRuntime::new()?;
1735    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1736    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?)?;
1737    ///     s.reduce_max(&x, None)
1738    /// })?;
1739    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[3.0]);
1740    /// # Ok::<(), tenferro_ad::Error>(())
1741    /// ```
1742    /// # Errors
1743    /// Returns a typed foreign-runtime, invalid-axis, unsupported-dtype, or backend error.
1744    pub fn reduce_max(
1745        &mut self,
1746        input: &EagerTensor,
1747        axes: Option<&[usize]>,
1748    ) -> Result<EagerTensor> {
1749        self.ensure_runtime(input)?;
1750        let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1751        crate::eager_ops::validate_eager_axes(
1752            "EagerSession::reduce_max",
1753            input.shape().len(),
1754            &axes,
1755        )?;
1756        self.run_unary(input, StdTensorOp::ReduceMax { axes })
1757    }
1758
1759    /// Reduce the minimum over selected axes in this borrowed session.
1760    /// `None` reduces every axis.
1761    ///
1762    /// # Examples
1763    /// ```rust
1764    /// use tenferro_ad::{EagerRuntime, Tensor};
1765    /// let ctx = EagerRuntime::new()?;
1766    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1767    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?)?;
1768    ///     s.reduce_min(&x, None)
1769    /// })?;
1770    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[2.0]);
1771    /// # Ok::<(), tenferro_ad::Error>(())
1772    /// ```
1773    /// # Errors
1774    /// Returns a typed foreign-runtime, invalid-axis, unsupported-dtype, or backend error.
1775    pub fn reduce_min(
1776        &mut self,
1777        input: &EagerTensor,
1778        axes: Option<&[usize]>,
1779    ) -> Result<EagerTensor> {
1780        self.ensure_runtime(input)?;
1781        let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1782        crate::eager_ops::validate_eager_axes(
1783            "EagerSession::reduce_min",
1784            input.shape().len(),
1785            &axes,
1786        )?;
1787        self.run_unary(input, StdTensorOp::ReduceMin { axes })
1788    }
1789
1790    /// Duplicate an eager value into an independent tensor within the caller's
1791    /// borrowed session, preserving its dtype and placement.
1792    ///
1793    /// # Examples
1794    ///
1795    /// ```rust
1796    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1797    /// let ctx = EagerRuntime::new()?;
1798    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?, ctx.clone())?;
1799    /// let copy = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.duplicate_value(&x))?;
1800    /// assert_eq!(copy.as_slice::<f64>()?, &[2.0]);
1801    /// # Ok::<(), tenferro_ad::Error>(())
1802    /// ```
1803    ///
1804    /// # Errors
1805    ///
1806    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1807    /// runtime/backend error when the retained value cannot be duplicated.
1808    pub fn duplicate_value(&mut self, input: &EagerTensor) -> Result<Tensor> {
1809        self.ensure_runtime(input)?;
1810        input.duplicate_value_in_session(&mut *self.engine)
1811    }
1812
1813    /// Import an untracked leaf within this borrowed session.
1814    ///
1815    /// `tensor` must already be usable by this runtime's backend: a host
1816    /// tensor on a CPU runtime, or a tensor already on the device of a CUDA
1817    /// or WebGPU runtime. No host/device transfer happens here. To import host
1818    /// data into a device runtime, use [`Self::constant_from_host`], which
1819    /// uploads first; on a CPU runtime the two are equivalent.
1820    ///
1821    /// # Examples
1822    ///
1823    /// ```rust
1824    /// use tenferro_ad::{EagerRuntime, Tensor};
1825    /// let ctx = EagerRuntime::new()?;
1826    /// let c = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1827    ///     session.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)
1828    /// })?;
1829    /// assert_eq!(c.value()?.as_slice::<f64>()?, &[2.0]);
1830    /// # Ok::<(), tenferro_ad::Error>(())
1831    /// ```
1832    ///
1833    /// # Errors
1834    ///
1835    /// Returns [`Error::TensorRuntime`] for a typed backend failure when the value cannot be
1836    /// registered in the session, or [`Error::RuntimeState`] when the runtime's
1837    /// value registry is unavailable.
1838    pub fn constant_from(&mut self, tensor: Tensor) -> Result<EagerTensor> {
1839        EagerTensor::new_leaf_in_session(Arc::clone(self.runtime), tensor, false, &mut *self.engine)
1840    }
1841
1842    /// Upload a host tensor and import it as an untracked leaf in this session.
1843    ///
1844    /// Unlike [`Self::constant_from`], this explicitly crosses the host/device
1845    /// boundary: the host `tensor` is uploaded to this runtime's backend
1846    /// (a host copy on a CPU runtime) and the uploaded value becomes the leaf.
1847    /// Use it whenever the source data lives on the host and the runtime may
1848    /// be a device runtime; use [`Self::constant_from`] for a tensor that is
1849    /// already resident on the backend.
1850    ///
1851    /// # Examples
1852    /// ```rust
1853    /// use tenferro_ad::{EagerRuntime, Tensor};
1854    /// let ctx = EagerRuntime::new()?;
1855    /// let c = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1856    ///     s.constant_from_host(Tensor::from_vec_col_major([1], vec![2.0_f64])?)
1857    /// })?;
1858    /// assert_eq!(c.value()?.as_slice::<f64>()?, &[2.0]);
1859    /// # Ok::<(), tenferro_ad::Error>(())
1860    /// ```
1861    /// # Errors
1862    /// Returns [`Error::TensorRuntime`] for a typed backend failure, including a host-tensor
1863    /// upload failure, or [`Error::RuntimeState`] when the runtime's value
1864    /// registry is unavailable.
1865    pub fn constant_from_host(&mut self, tensor: Tensor) -> Result<EagerTensor> {
1866        let uploaded = self
1867            .engine
1868            .upload_host_tensor(TensorRead::from_tensor(&tensor))
1869            .map_err(Error::from)?;
1870        self.constant_from(uploaded)
1871    }
1872
1873    /// Import a trainable leaf within this borrowed session.
1874    ///
1875    /// # Examples
1876    ///
1877    /// ```rust
1878    /// use tenferro_ad::{EagerRuntime, Tensor};
1879    /// let ctx = EagerRuntime::new()?;
1880    /// let x = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1881    ///     session.variable_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)
1882    /// })?;
1883    /// assert!(x.tracks_grad());
1884    /// # Ok::<(), tenferro_ad::Error>(())
1885    /// ```
1886    ///
1887    /// # Errors
1888    ///
1889    /// Returns [`Error::TensorRuntime`] for a typed backend failure when the value cannot be
1890    /// registered, or [`Error::RuntimeState`] when the runtime's value or
1891    /// gradient registry is unavailable.
1892    pub fn variable_from(&mut self, tensor: Tensor) -> Result<EagerTensor> {
1893        EagerTensor::new_leaf_in_session(Arc::clone(self.runtime), tensor, true, &mut *self.engine)
1894    }
1895
1896    /// Add eager tensors with the same broadcast and AD rules as the eager
1897    /// operation surface, reusing this borrowed execution session.
1898    ///
1899    /// # Examples
1900    ///
1901    /// ```rust
1902    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1903    /// let ctx = EagerRuntime::new()?;
1904    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1905    /// let scalar = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?, ctx.clone())?;
1906    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.add(&x, &scalar))?;
1907    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[4.0, 5.0]);
1908    /// # Ok::<(), tenferro_ad::Error>(())
1909    /// ```
1910    ///
1911    /// # Errors
1912    ///
1913    /// Returns [`Error::ContextMismatch`] for a tensor from another runtime,
1914    /// or a typed broadcast/backend error for the operands.
1915    pub fn add(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1916        self.run_binary("add", lhs, rhs, StdTensorOp::Add)
1917    }
1918
1919    /// Subtract eager tensors within this borrowed session.
1920    ///
1921    /// # Examples
1922    ///
1923    /// ```rust
1924    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1925    /// let ctx = EagerRuntime::new()?;
1926    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?, ctx.clone())?;
1927    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.sub(&x, &x))?;
1928    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1929    /// # Ok::<(), tenferro_ad::Error>(())
1930    /// ```
1931    ///
1932    /// # Errors
1933    ///
1934    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1935    /// broadcast/backend error for the operands.
1936    pub fn sub(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1937        self.run_binary("sub", lhs, rhs, StdTensorOp::Sub)
1938    }
1939
1940    /// Multiply eager tensors within this borrowed session.
1941    ///
1942    /// # Examples
1943    ///
1944    /// ```rust
1945    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1946    /// let ctx = EagerRuntime::new()?;
1947    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?, ctx.clone())?;
1948    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.mul(&x, &x))?;
1949    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[9.0]);
1950    /// # Ok::<(), tenferro_ad::Error>(())
1951    /// ```
1952    ///
1953    /// # Errors
1954    ///
1955    /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1956    /// broadcast/backend error for the operands.
1957    pub fn mul(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1958        self.run_binary("mul", lhs, rhs, StdTensorOp::Mul)
1959    }
1960
1961    /// Divide eager tensors elementwise with broadcast rules.
1962    ///
1963    /// # Examples
1964    /// ```rust
1965    /// use tenferro_ad::{EagerRuntime, Tensor};
1966    /// let ctx = EagerRuntime::new()?;
1967    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1968    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![6.0_f64])?)?;
1969    ///     let divisor = s.constant_from(Tensor::from_vec_col_major(vec![], vec![2.0_f64])?)?;
1970    ///     s.div(&x, &divisor)
1971    /// })?;
1972    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[3.0]);
1973    /// # Ok::<(), tenferro_ad::Error>(())
1974    /// ```
1975    /// # Errors
1976    /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
1977    /// runtime, a validation error with
1978    /// `ValidationError::ShapeMismatch` when the operands cannot broadcast, or
1979    /// [`Error::TensorRuntime`] for a typed backend failure (including integer division by
1980    /// zero).
1981    pub fn div(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1982        self.run_binary("div", lhs, rhs, StdTensorOp::Div)
1983    }
1984
1985    /// Compute the elementwise remainder with broadcast rules.
1986    ///
1987    /// # Examples
1988    /// ```rust
1989    /// use tenferro_ad::{EagerRuntime, Tensor};
1990    /// let ctx = EagerRuntime::new()?;
1991    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1992    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![5.0_f64])?)?;
1993    ///     let divisor = s.constant_from(Tensor::from_vec_col_major(vec![], vec![2.0_f64])?)?;
1994    ///     s.rem(&x, &divisor)
1995    /// })?;
1996    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0]);
1997    /// # Ok::<(), tenferro_ad::Error>(())
1998    /// ```
1999    /// # Errors
2000    /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
2001    /// runtime, a validation error with
2002    /// `ValidationError::ShapeMismatch` when the operands cannot broadcast, or
2003    /// [`Error::TensorRuntime`] for a typed backend failure (including an integer remainder by
2004    /// zero).
2005    pub fn rem(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2006        self.run_binary("rem", lhs, rhs, StdTensorOp::Rem)
2007    }
2008
2009    /// Raise eager tensor elements to broadcast exponents.
2010    ///
2011    /// # Examples
2012    /// ```rust
2013    /// use tenferro_ad::{EagerRuntime, Tensor};
2014    /// let ctx = EagerRuntime::new()?;
2015    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2016    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2017    ///     let exponent = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2018    ///     s.pow(&x, &exponent)
2019    /// })?;
2020    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[8.0]);
2021    /// # Ok::<(), tenferro_ad::Error>(())
2022    /// ```
2023    /// # Errors
2024    /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
2025    /// runtime, a validation error with
2026    /// `ValidationError::ShapeMismatch` when the operands cannot broadcast, or
2027    /// [`Error::TensorRuntime`] for a typed backend failure (including a negative integer
2028    /// exponent).
2029    pub fn pow(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2030        self.run_binary("pow", lhs, rhs, StdTensorOp::Pow)
2031    }
2032
2033    /// Compute the elementwise maximum under broadcast rules.
2034    ///
2035    /// # Examples
2036    /// ```rust
2037    /// use tenferro_ad::{EagerRuntime, Tensor};
2038    /// let ctx = EagerRuntime::new()?;
2039    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2040    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2041    ///     let bound = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2042    ///     s.maximum(&x, &bound)
2043    /// })?;
2044    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[3.0]);
2045    /// # Ok::<(), tenferro_ad::Error>(())
2046    /// ```
2047    /// # Errors
2048    /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2049    pub fn maximum(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2050        self.run_binary("maximum", lhs, rhs, StdTensorOp::Maximum)
2051    }
2052
2053    /// Compute the elementwise minimum under broadcast rules.
2054    ///
2055    /// # Examples
2056    /// ```rust
2057    /// use tenferro_ad::{EagerRuntime, Tensor};
2058    /// let ctx = EagerRuntime::new()?;
2059    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2060    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2061    ///     let bound = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2062    ///     s.minimum(&x, &bound)
2063    /// })?;
2064    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
2065    /// # Ok::<(), tenferro_ad::Error>(())
2066    /// ```
2067    /// # Errors
2068    /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2069    pub fn minimum(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2070        self.run_binary("minimum", lhs, rhs, StdTensorOp::Minimum)
2071    }
2072
2073    /// Compare eager tensors elementwise under broadcast rules.
2074    ///
2075    /// # Examples
2076    /// ```rust
2077    /// use tenferro_ad::{CompareDir, EagerRuntime, Tensor};
2078    /// let ctx = EagerRuntime::new()?;
2079    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2080    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2081    ///     let bound = s.constant_from(Tensor::from_vec_col_major(vec![], vec![1.0_f64])?)?;
2082    ///     s.compare(&x, &bound, CompareDir::Gt)
2083    /// })?;
2084    /// assert_eq!(y.value()?.as_slice::<bool>()?, &[true]);
2085    /// # Ok::<(), tenferro_ad::Error>(())
2086    /// ```
2087    /// # Errors
2088    /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2089    pub fn compare(
2090        &mut self,
2091        lhs: &EagerTensor,
2092        rhs: &EagerTensor,
2093        dir: CompareDir,
2094    ) -> Result<EagerTensor> {
2095        self.run_binary("compare", lhs, rhs, StdTensorOp::Compare(dir))
2096    }
2097
2098    /// Select eager values elementwise using a broadcast boolean condition.
2099    ///
2100    /// # Examples
2101    /// ```rust
2102    /// use tenferro_ad::{EagerRuntime, Tensor};
2103    /// let ctx = EagerRuntime::new()?;
2104    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2105    ///     let condition = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![true, false])?)?;
2106    ///     let yes = s.constant_from(Tensor::from_vec_col_major(vec![], vec![10.0_f64])?)?;
2107    ///     let no = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
2108    ///     s.where_select(&condition, &yes, &no)
2109    /// })?;
2110    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[10.0, 2.0]);
2111    /// # Ok::<(), tenferro_ad::Error>(())
2112    /// ```
2113    /// # Errors
2114    /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2115    pub fn where_select(
2116        &mut self,
2117        condition: &EagerTensor,
2118        on_true: &EagerTensor,
2119        on_false: &EagerTensor,
2120    ) -> Result<EagerTensor> {
2121        self.run_ternary(
2122            "where_select",
2123            condition,
2124            on_true,
2125            on_false,
2126            StdTensorOp::Select,
2127        )
2128    }
2129
2130    /// Alias for [`Self::where_select`] with the same borrowed-session semantics.
2131    ///
2132    /// # Examples
2133    /// ```rust
2134    /// use tenferro_ad::{EagerRuntime, Tensor};
2135    /// let ctx = EagerRuntime::new()?;
2136    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2137    ///     let predicate = s.constant_from(Tensor::from_vec_col_major(vec![], vec![true])?)?;
2138    ///     let yes = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2139    ///     let no = s.constant_from(Tensor::from_vec_col_major(vec![], vec![4.0_f64])?)?;
2140    ///     s.select(&predicate, &yes, &no)
2141    /// })?;
2142    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[3.0]);
2143    /// # Ok::<(), tenferro_ad::Error>(())
2144    /// ```
2145    /// # Errors
2146    /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2147    pub fn select(
2148        &mut self,
2149        condition: &EagerTensor,
2150        on_true: &EagerTensor,
2151        on_false: &EagerTensor,
2152    ) -> Result<EagerTensor> {
2153        self.where_select(condition, on_true, on_false)
2154    }
2155
2156    /// Clamp eager values elementwise between broadcast lower and upper bounds.
2157    ///
2158    /// # Examples
2159    /// ```rust
2160    /// use tenferro_ad::{EagerRuntime, Tensor};
2161    /// let ctx = EagerRuntime::new()?;
2162    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2163    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![-2.0_f64, 5.0])?)?;
2164    ///     let lo = s.constant_from(Tensor::from_vec_col_major(vec![], vec![-1.0_f64])?)?;
2165    ///     let hi = s.constant_from(Tensor::from_vec_col_major(vec![], vec![4.0_f64])?)?;
2166    ///     s.clamp(&x, &lo, &hi)
2167    /// })?;
2168    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-1.0, 4.0]);
2169    /// # Ok::<(), tenferro_ad::Error>(())
2170    /// ```
2171    /// # Errors
2172    /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2173    pub fn clamp(
2174        &mut self,
2175        input: &EagerTensor,
2176        lower: &EagerTensor,
2177        upper: &EagerTensor,
2178    ) -> Result<EagerTensor> {
2179        self.run_ternary("clamp", input, lower, upper, StdTensorOp::Clamp)
2180    }
2181
2182    /// Contract eager tensors according to a dot-general dimension mapping.
2183    ///
2184    /// The output layout is `[lhs free..., rhs free..., batch...]`: batch axes
2185    /// come last (see [`DotGeneralConfig`](tenferro_runtime::DotGeneralConfig)).
2186    ///
2187    /// # Examples
2188    ///
2189    /// ```rust
2190    /// use tenferro_ad::{DotGeneralConfig, EagerRuntime, Tensor};
2191    /// let ctx = EagerRuntime::new()?;
2192    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
2193    ///     let lhs = session.variable_from(Tensor::from_vec_col_major(vec![1, 2], vec![2.0_f64, 3.0])?)?;
2194    ///     let rhs = session.constant_from(Tensor::from_vec_col_major(vec![2, 1], vec![4.0_f64, 5.0])?)?;
2195    ///     session.dot_general(&lhs, &rhs, DotGeneralConfig {
2196    ///         lhs_contracting_dims: [1].as_slice().into(),
2197    ///         rhs_contracting_dims: [0].as_slice().into(),
2198    ///         lhs_batch_dims: [].as_slice().into(),
2199    ///         rhs_batch_dims: [].as_slice().into(),
2200    ///     })
2201    /// })?;
2202    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[23.0]);
2203    /// # Ok::<(), tenferro_ad::Error>(())
2204    /// ```
2205    ///
2206    /// # Errors
2207    ///
2208    /// Returns [`Error::ContextMismatch`] for a foreign eager runtime,
2209    /// a typed validation error for incompatible contraction dimensions,
2210    /// or the backend's typed execution error.
2211    pub fn dot_general(
2212        &mut self,
2213        lhs: &EagerTensor,
2214        rhs: &EagerTensor,
2215        config: DotGeneralConfig,
2216    ) -> Result<EagerTensor> {
2217        self.ensure_runtime(lhs)?;
2218        self.ensure_runtime(rhs)?;
2219        config
2220            .validate_dims_with_ranks(lhs.shape().len(), rhs.shape().len())
2221            .map_err(Error::TensorRuntime)?;
2222        EagerTensor::nary_op_in_session(
2223            &[lhs, rhs],
2224            StdTensorOp::DotGeneral { config },
2225            &mut *self.engine,
2226        )
2227    }
2228
2229    /// Scale an eager tensor by a real scalar in this borrowed session.
2230    /// Integer factors are rounded; finite zero maps to `false` for boolean inputs.
2231    ///
2232    /// # Examples
2233    /// ```rust
2234    /// use tenferro_ad::{EagerRuntime, Tensor};
2235    /// let ctx = EagerRuntime::new()?;
2236    /// let scaled = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2237    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
2238    ///     s.scale_real(&x, 2.0)
2239    /// })?;
2240    /// assert_eq!(scaled.value()?.as_slice::<f64>()?, &[2.0, 4.0]);
2241    /// # Ok::<(), tenferro_ad::Error>(())
2242    /// ```
2243    /// # Errors
2244    /// Returns a typed foreign-runtime, invalid-factor/dtype, or backend error.
2245    pub fn scale_real(&mut self, input: &EagerTensor, factor: f64) -> Result<EagerTensor> {
2246        self.ensure_runtime(input)?;
2247        let scalar = tenferro_runtime::scale::real_scale_scalar(input.dtype(), factor)?;
2248        let scalar = self.constant_from_host(scalar)?;
2249        self.mul(input, &scalar)
2250    }
2251
2252    /// Scale a complex eager tensor by a complex scalar in this borrowed session.
2253    ///
2254    /// # Examples
2255    /// ```rust
2256    /// use num_complex::Complex64;
2257    /// use tenferro_ad::{EagerRuntime, Tensor};
2258    /// let ctx = EagerRuntime::new()?;
2259    /// let scaled = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2260    ///     let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![Complex64::new(1.0, 2.0)])?)?;
2261    ///     s.scale_complex(&x, Complex64::new(0.0, 1.0))
2262    /// })?;
2263    /// assert_eq!(scaled.value()?.as_slice::<Complex64>()?, &[Complex64::new(-2.0, 1.0)]);
2264    /// # Ok::<(), tenferro_ad::Error>(())
2265    /// ```
2266    /// # Errors
2267    /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
2268    /// runtime, [`Error::TensorRuntime`] containing
2269    /// `ValidationError::InvalidArgument` when the input dtype is not complex,
2270    /// or [`Error::TensorRuntime`] for a typed backend failure.
2271    pub fn scale_complex(&mut self, input: &EagerTensor, factor: Complex64) -> Result<EagerTensor> {
2272        self.ensure_runtime(input)?;
2273        let scalar = tenferro_runtime::scale::complex_scale_scalar(input.dtype(), factor)?;
2274        let scalar = self.constant_from_host(scalar)?;
2275        self.mul(input, &scalar)
2276    }
2277
2278    /// Multiply two rank-2 eager tensors in this borrowed session.
2279    ///
2280    /// # Examples
2281    /// ```rust
2282    /// use tenferro_ad::{EagerRuntime, Tensor};
2283    /// let ctx = EagerRuntime::new()?;
2284    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2285    ///     let a = s.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![2.0_f64])?)?;
2286    ///     let b = s.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![3.0_f64])?)?;
2287    ///     s.matmul(&a, &b)
2288    /// })?;
2289    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[6.0]);
2290    /// # Ok::<(), tenferro_ad::Error>(())
2291    /// ```
2292    /// # Errors
2293    /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
2294    /// runtime, a validation error with
2295    /// `ValidationError::RankMismatch` or `ValidationError::ShapeMismatch` when
2296    /// the operands are not rank-2 with matching inner dimensions, a dtype
2297    /// mismatch between the operands, or [`Error::TensorRuntime`] for a typed backend failure.
2298    pub fn matmul(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2299        self.ensure_runtime(lhs)?;
2300        self.ensure_runtime(rhs)?;
2301        let lhs_shape = lhs.shape();
2302        let rhs_shape = rhs.shape();
2303        if lhs_shape.len() != 2 {
2304            return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, lhs_shape.len()).into());
2305        }
2306        if rhs_shape.len() != 2 {
2307            return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, rhs_shape.len()).into());
2308        }
2309        if lhs_shape[1] != rhs_shape[0] {
2310            return Err(
2311                tenferro_tensor::Error::shape_mismatch("matmul", lhs_shape, rhs_shape).into(),
2312            );
2313        }
2314        self.dot_general(
2315            lhs,
2316            rhs,
2317            DotGeneralConfig {
2318                lhs_contracting_dims: [1].as_slice().into(),
2319                rhs_contracting_dims: [0].as_slice().into(),
2320                lhs_batch_dims: [].as_slice().into(),
2321                rhs_batch_dims: [].as_slice().into(),
2322            },
2323        )
2324    }
2325
2326    /// Contract eagerly with optional conjugation of either operand.
2327    /// Untracked operands use the backend's conjugating contraction directly;
2328    /// tracked operands record explicit conjugations for reverse-mode AD.
2329    ///
2330    /// # Examples
2331    ///
2332    /// ```rust
2333    /// use tenferro_ad::{DotGeneralConfig, EagerRuntime, Tensor};
2334    /// let ctx = EagerRuntime::new()?;
2335    /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
2336    ///     let lhs = session.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![2.0_f64])?)?;
2337    ///     let rhs = session.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![3.0_f64])?)?;
2338    ///     session.dot_general_with_conj(&lhs, &rhs, DotGeneralConfig {
2339    ///         lhs_contracting_dims: [1].as_slice().into(),
2340    ///         rhs_contracting_dims: [0].as_slice().into(),
2341    ///         lhs_batch_dims: [].as_slice().into(),
2342    ///         rhs_batch_dims: [].as_slice().into(),
2343    ///     }, true, false)
2344    /// })?;
2345    /// assert_eq!(result.value()?.as_slice::<f64>()?, &[6.0]);
2346    /// # Ok::<(), tenferro_ad::Error>(())
2347    /// ```
2348    ///
2349    /// # Errors
2350    ///
2351    /// Returns [`Error::ContextMismatch`] for a foreign eager runtime,
2352    /// a typed validation error for invalid dimensions, or a backend error.
2353    pub fn dot_general_with_conj(
2354        &mut self,
2355        lhs: &EagerTensor,
2356        rhs: &EagerTensor,
2357        config: DotGeneralConfig,
2358        lhs_conj: bool,
2359        rhs_conj: bool,
2360    ) -> Result<EagerTensor> {
2361        self.ensure_runtime(lhs)?;
2362        self.ensure_runtime(rhs)?;
2363        config
2364            .validate_dims_with_ranks(lhs.shape().len(), rhs.shape().len())
2365            .map_err(Error::TensorRuntime)?;
2366        if !lhs.requires_grad && !rhs.requires_grad {
2367            let output = crate::eager_exec::exec_dot_general_with_conj_on_tensor_reads_in_session(
2368                lhs.tensor_read(),
2369                rhs.tensor_read(),
2370                &config,
2371                lhs_conj,
2372                rhs_conj,
2373                &mut *self.engine,
2374            )?;
2375            return EagerTensor::new_untracked_result(Arc::clone(self.runtime), output);
2376        }
2377        let lhs = if lhs_conj {
2378            self.conj(lhs)?
2379        } else {
2380            lhs.clone()
2381        };
2382        let rhs = if rhs_conj {
2383            self.conj(rhs)?
2384        } else {
2385            rhs.clone()
2386        };
2387        self.dot_general(&lhs, &rhs, config)
2388    }
2389
2390    fn run_binary(
2391        &mut self,
2392        name: &'static str,
2393        lhs: &EagerTensor,
2394        rhs: &EagerTensor,
2395        op: StdTensorOp,
2396    ) -> Result<EagerTensor> {
2397        self.ensure_runtime(lhs)?;
2398        self.ensure_runtime(rhs)?;
2399        let (lhs, rhs) = crate::eager_ops::broadcast_binary_in_session(name, lhs, rhs, self)?;
2400        EagerTensor::nary_op_in_session(&[lhs.as_ref(), rhs.as_ref()], op, &mut *self.engine)
2401    }
2402
2403    fn run_ternary(
2404        &mut self,
2405        name: &'static str,
2406        first: &EagerTensor,
2407        second: &EagerTensor,
2408        third: &EagerTensor,
2409        op: StdTensorOp,
2410    ) -> Result<EagerTensor> {
2411        self.ensure_runtime(first)?;
2412        self.ensure_runtime(second)?;
2413        self.ensure_runtime(third)?;
2414        let (first, second, third) =
2415            crate::eager_ops::broadcast_ternary_in_session(name, first, second, third, self)?;
2416        EagerTensor::nary_op_in_session(&[&first, &second, &third], op, &mut *self.engine)
2417    }
2418
2419    /// Apply one standard tensor op in this borrowed session and record it
2420    /// for AD when needed.
2421    ///
2422    /// Extension crates use this when an extension-level eager operation
2423    /// expands into ordinary `StdTensorOp` nodes instead of a custom extension
2424    /// primitive: all of them run in this one backend session.
2425    ///
2426    /// # Examples
2427    ///
2428    /// ```rust
2429    /// use tenferro_ad::{EagerRuntime, Tensor};
2430    /// use tenferro_ops::std_tensor_op::StdTensorOp;
2431    ///
2432    /// let ctx = EagerRuntime::new()?;
2433    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2434    ///     let x = s.variable_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
2435    ///     let negated = s.apply_standard_op(StdTensorOp::Neg, &[&x])?;
2436    ///     s.apply_standard_op(StdTensorOp::Mul, &[&negated, &x])
2437    /// })?;
2438    /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-1.0, -4.0]);
2439    /// assert!(y.tracks_grad());
2440    /// # Ok::<(), tenferro_ad::Error>(())
2441    /// ```
2442    ///
2443    /// # Errors
2444    ///
2445    /// Returns [`Error::TensorRuntime`] containing
2446    /// [`tenferro_tensor::ValidationError::InvalidArgument`] for an extension
2447    /// op, [`Error::ContextMismatch`] for a tensor from another runtime, a
2448    /// typed input-count error, or the backend's typed execution error.
2449    pub fn apply_standard_op(
2450        &mut self,
2451        op: StdTensorOp,
2452        inputs: &[&EagerTensor],
2453    ) -> Result<EagerTensor> {
2454        if matches!(op, StdTensorOp::Extension(_)) {
2455            return Err(Error::invalid_argument(
2456                "EagerSession::apply_standard_op",
2457                ErrorPhase::Execution,
2458                "op",
2459                "Extension ops must be passed to apply_eager",
2460            ));
2461        }
2462        for input in inputs {
2463            self.ensure_runtime(input)?;
2464        }
2465        EagerTensor::nary_op_in_session(inputs, op, &mut *self.engine)
2466    }
2467
2468    /// Borrow the backend session this eager session runs on.
2469    ///
2470    /// Extension crates use it to run their backend kernels on untracked
2471    /// values inside the same execution region instead of entering a second
2472    /// session, which would be rejected as reentry. It grants the same access
2473    /// as passing the engine directly.
2474    ///
2475    /// # Examples
2476    ///
2477    /// ```rust
2478    /// use tenferro_ad::{EagerRuntime, Tensor};
2479    /// use tenferro_tensor::TensorRead;
2480    ///
2481    /// let ctx = EagerRuntime::new()?;
2482    /// let x = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0])?;
2483    /// let copy = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2484    ///     s.backend_session()
2485    ///         .to_contiguous_read(TensorRead::from_tensor(&x))
2486    ///         .map_err(tenferro_ad::Error::from)
2487    /// })?;
2488    /// assert_eq!(copy.as_slice::<f64>()?, &[1.0, -2.0]);
2489    /// # Ok::<(), Box<dyn std::error::Error>>(())
2490    /// ```
2491    pub fn backend_session(&mut self) -> &mut dyn Engine {
2492        &mut *self.engine
2493    }
2494
2495    /// Borrow the engine this eager session runs on.
2496    ///
2497    /// # Examples
2498    ///
2499    /// ```rust
2500    /// use tenferro_ad::EagerRuntime;
2501    /// use tenferro_cpu::CpuEngine;
2502    ///
2503    /// let ctx = EagerRuntime::new()?;
2504    /// let mut engine = CpuEngine::serial();
2505    /// let id = ctx.with_eager_session(&mut engine, |s| {
2506    ///     Ok::<_, tenferro_ad::Error>(s.engine().target().engine_id().clone())
2507    /// })?;
2508    /// assert_eq!(id, tenferro_cpu::runtime_engine_id()?);
2509    /// # Ok::<(), Box<dyn std::error::Error>>(())
2510    /// ```
2511    pub fn engine(&mut self) -> &mut dyn ExecutionEngine {
2512        &mut *self.engine
2513    }
2514
2515    fn run_unary(&mut self, input: &EagerTensor, op: StdTensorOp) -> Result<EagerTensor> {
2516        self.ensure_runtime(input)?;
2517        EagerTensor::nary_op_in_session(&[input], op, &mut *self.engine)
2518    }
2519
2520    pub(crate) fn record_outputs(
2521        &mut self,
2522        op: &StdTensorOp,
2523        outputs: &[&Tensor],
2524        inputs: &[&EagerTensor],
2525    ) -> Result<RecordedEagerOutputs> {
2526        record_eager_outputs_in_session(op, outputs, inputs, &mut *self.engine)
2527    }
2528
2529    pub(crate) fn ensure_runtime(&self, input: &EagerTensor) -> Result<()> {
2530        if !Arc::ptr_eq(self.runtime, &input.ctx) {
2531            return Err(Error::ContextMismatch {
2532                lhs: self.runtime.id(),
2533                rhs: input.ctx_id(),
2534            });
2535        }
2536        Ok(())
2537    }
2538
2539    pub(crate) fn runtime(&self) -> &Arc<EagerRuntime> {
2540        self.runtime
2541    }
2542
2543    /// Run `f` on this runtime's extension cache store from inside the
2544    /// session.
2545    ///
2546    /// Operation families use this for their own prepared-plan caches without
2547    /// reopening the runtime: the eager owner is already locked, and the cache
2548    /// lock is taken second, as in every extension execution region.
2549    ///
2550    /// # Examples
2551    ///
2552    /// ```rust
2553    /// use tenferro_ad::EagerRuntime;
2554    ///
2555    /// let ctx = EagerRuntime::new()?;
2556    /// let entries = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
2557    ///     session.with_extension_caches(|caches| caches.len())
2558    /// })?;
2559    /// assert_eq!(entries, 0);
2560    /// # Ok::<(), tenferro_ad::Error>(())
2561    /// ```
2562    ///
2563    /// # Errors
2564    ///
2565    /// Returns a runtime-state error when the extension cache lock is
2566    /// poisoned.
2567    pub fn with_extension_caches<R>(
2568        &mut self,
2569        f: impl FnOnce(&mut tenferro_runtime::ExtensionCacheStore) -> R,
2570    ) -> Result<R> {
2571        let result = with_lent_extension_caches(self.engine, |_, caches| f(caches));
2572        Ok(result)
2573    }
2574}
2575
2576impl fmt::Debug for EagerRuntime {
2577    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2578        let mut debug = f.debug_struct("EagerRuntime");
2579        debug.field("id", &self.id);
2580        debug.field("runtime_id", &self.runtime.id());
2581        debug.field("runtime_epoch", &self.runtime.epoch().ok());
2582        match self.extension_install_lock.try_lock() {
2583            Ok(_) => {
2584                debug.field("extension_install_lock", &"<unlocked>");
2585            }
2586            Err(_) => {
2587                debug.field("extension_install_lock", &"<locked>");
2588            }
2589        }
2590        debug.field("semantic_extension_rules", &self.semantic_extension_rules);
2591        match self.grad_slots.try_lock() {
2592            Ok(slots) => {
2593                debug.field("grad_slots_len", &slots.len());
2594            }
2595            Err(_) => {
2596                debug.field("grad_slots_len", &"<locked>");
2597            }
2598        }
2599        match self.value_records.try_lock() {
2600            Ok(records) => {
2601                debug.field("value_records_len", &records.len());
2602            }
2603            Err(_) => {
2604                debug.field("value_records_len", &"<locked>");
2605            }
2606        }
2607        match self.ad_transform_cache.stats() {
2608            Ok(stats) => {
2609                debug.field("ad_transform_cache_stats", &stats);
2610            }
2611            Err(err) => {
2612                debug.field("ad_transform_cache_stats", &format_args!("{err}"));
2613            }
2614        }
2615        match self.prepared_derivative_cache.try_lock() {
2616            Ok(cache) => {
2617                debug.field("prepared_derivative_cache_stats", &cache.stats());
2618            }
2619            Err(_) => {
2620                debug.field("prepared_derivative_cache_stats", &"<locked>");
2621            }
2622        }
2623        debug.finish_non_exhaustive()
2624    }
2625}
2626
2627impl EagerRuntime {
2628    fn lock_extension_install(&self) -> Result<MutexGuard<'_, ()>> {
2629        self.extension_install_lock.lock().map_err(|_| {
2630            Error::runtime_state(
2631                "eager_extension_install",
2632                ErrorPhase::Execution,
2633                "lock poisoned",
2634            )
2635        })
2636    }
2637
2638    fn lock_prepared_derivative_cache(&self) -> Result<MutexGuard<'_, PreparedDerivativeCache>> {
2639        self.prepared_derivative_cache.lock().map_err(|_| {
2640            Error::runtime_state(
2641                "prepared_derivative_cache",
2642                ErrorPhase::Execution,
2643                "lock poisoned",
2644            )
2645        })
2646    }
2647
2648    fn lock_grad_slots(
2649        &self,
2650    ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
2651        self.grad_slots.lock().map_err(|_| {
2652            Error::runtime_state(
2653                "eager_gradient_slots",
2654                ErrorPhase::Execution,
2655                "lock poisoned",
2656            )
2657        })
2658    }
2659
2660    fn lock_value_records(
2661        &self,
2662    ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>> {
2663        self.value_records.lock().map_err(|_| {
2664            Error::runtime_state(
2665                "eager_value_registry",
2666                ErrorPhase::Execution,
2667                "lock poisoned",
2668            )
2669        })
2670    }
2671
2672    fn with_rules_and_cache(
2673        semantic_extension_rules: SemanticExtensionRuleSet,
2674        ad_transform_cache: Arc<AdTransformCache>,
2675    ) -> Result<Self> {
2676        let runtime = Runtime::builder()
2677            .build()
2678            .map_err(|source| runtime_config_error("EagerRuntime::new", source))?;
2679        Ok(Self {
2680            id: ContextId::fresh(),
2681            runtime,
2682            extension_install_lock: Mutex::new(()),
2683            semantic_extension_rules,
2684            grad_slots: Mutex::new(HashMap::new()),
2685            value_records: Mutex::new(HashMap::new()),
2686            ad_transform_cache,
2687            prepared_derivative_cache: Mutex::new(PreparedDerivativeCache::default()),
2688        })
2689    }
2690
2691    /// Create a shared eager execution context.
2692    ///
2693    /// The context stores no engine: every operation runs on the engine the
2694    /// caller passes in.
2695    ///
2696    /// # Examples
2697    ///
2698    /// ```
2699    /// use tenferro_ad::EagerRuntime;
2700    ///
2701    /// let ctx = EagerRuntime::new()?;
2702    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
2703    /// # Ok::<(), tenferro_ad::Error>(())
2704    /// ```
2705    ///
2706    /// # Errors
2707    ///
2708    /// Returns [`Error::RuntimeStateSource`] when the prepared-program runtime
2709    /// cannot be configured, preserving the underlying [`RuntimeConfigError`].
2710    pub fn new() -> Result<Arc<Self>> {
2711        Ok(Arc::new(Self::with_rules_and_cache(
2712            SemanticExtensionRuleSet::default(),
2713            Arc::new(AdTransformCache::new()),
2714        )?))
2715    }
2716
2717    /// Create a shared eager execution context that uses the semantic
2718    /// extension rules and AD transform cache of `ad`.
2719    ///
2720    /// # Examples
2721    ///
2722    /// ```
2723    /// use tenferro_ad::{AdContext, EagerRuntime};
2724    ///
2725    /// let ad = AdContext::builder().build().unwrap();
2726    /// let ctx = EagerRuntime::with_ad_context(&ad)?;
2727    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
2728    /// # Ok::<(), tenferro_ad::Error>(())
2729    /// ```
2730    ///
2731    /// # Errors
2732    ///
2733    /// Returns [`Error::RuntimeStateSource`] when the prepared-program runtime
2734    /// cannot be configured.
2735    pub fn with_ad_context(ad: &AdContext) -> Result<Arc<Self>> {
2736        Ok(Arc::new(Self::with_rules_and_cache(
2737            ad.semantic_extension_rules().clone(),
2738            ad.ad_transform_cache(),
2739        )?))
2740    }
2741
2742    /// Register the kind of `engine` with this runtime if it is not registered
2743    /// yet, and return its target.
2744    pub(crate) fn ensure_engine(&self, engine: &mut dyn ExecutionEngine) -> Result<EngineTarget> {
2745        let target = engine.target();
2746        let snapshot = self
2747            .runtime
2748            .snapshot()
2749            .map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
2750        if let Some(registered) = snapshot.engine(target.engine_id()) {
2751            if registered.context_identity() != target.context_identity() {
2752                return Err(Error::unsupported(
2753                    "EagerRuntime::ensure_engine",
2754                    ErrorPhase::Execution,
2755                    format!(
2756                        "engine id {:?} is registered for a different engine type",
2757                        target.engine_id()
2758                    ),
2759                ));
2760            }
2761            return Ok(target);
2762        }
2763        let registration = engine_kind_registration(engine)
2764            .ok_or_else(|| unregisterable_engine_error(target.engine_id()))?;
2765        let registration = registration
2766            .map_err(|source| runtime_config_error("EagerRuntime::ensure_engine", source))?;
2767        let _install_guard = self.lock_extension_install()?;
2768        let snapshot = self
2769            .runtime
2770            .snapshot()
2771            .map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
2772        if snapshot.engine(target.engine_id()).is_none() {
2773            self.runtime
2774                .reconfigure(|edit| {
2775                    edit.register_engine(registration)?;
2776                    Ok(())
2777                })
2778                .map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
2779        }
2780        Ok(target)
2781    }
2782
2783    /// Return an opaque identifier for this context.
2784    ///
2785    /// # Examples
2786    ///
2787    /// ```
2788    /// use tenferro_ad::{EagerRuntime};
2789    ///
2790    /// let ctx = EagerRuntime::new()?;
2791    /// assert_ne!(ctx.id(), EagerRuntime::new()?.id());
2792    /// # Ok::<(), tenferro_ad::Error>(())
2793    /// ```
2794    pub fn id(&self) -> ContextId {
2795        self.id
2796    }
2797
2798    /// Disable eager operation recording on the current thread until the guard is dropped.
2799    ///
2800    /// This is useful for optimizer updates, metric calculations, and other
2801    /// eager computations that should not become part of the AD tape.
2802    ///
2803    /// # Examples
2804    ///
2805    /// ```
2806    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2807    ///
2808    /// let ctx = EagerRuntime::new()?;
2809    /// let x = EagerTensor::requires_grad_in(
2810    ///     Tensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap(),
2811    ///     ctx.clone(),
2812    /// )?;
2813    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2814    ///     let _guard = ctx.no_grad();
2815    ///     s.mul(&x, &x)
2816    /// })?;
2817    /// assert!(!y.tracks_grad());
2818    /// # Ok::<(), tenferro_ad::Error>(())
2819    /// ```
2820    pub fn no_grad(&self) -> EagerNoGradGuard {
2821        EAGER_NO_GRAD_DEPTH.with(|depth| {
2822            depth.set(depth.get().saturating_add(1));
2823        });
2824        EagerNoGradGuard {
2825            active: true,
2826            _not_send: PhantomData,
2827        }
2828    }
2829
2830    /// Keep semantic-trace recording active for untracked intermediates.
2831    ///
2832    /// See [`EagerTraceCaptureGuard`] for the full contract and an example.
2833    pub fn capture_trace(&self) -> EagerTraceCaptureGuard {
2834        EAGER_CAPTURE_DEPTH.with(|depth| {
2835            depth.set(depth.get().saturating_add(1));
2836        });
2837        EagerTraceCaptureGuard {
2838            active: true,
2839            _not_send: PhantomData,
2840        }
2841    }
2842
2843    /// Install or replace one extension module on this eager context's runtime.
2844    ///
2845    /// Eager extension wrappers call this as an idempotent "ensure installed"
2846    /// step. When the exact module instance (same module ID and allocation) is
2847    /// already installed, this is a read-only no-op that returns the current
2848    /// runtime epoch without acquiring the install lock or reconfiguring. The
2849    /// cold or replacement paths keep the transactional install-or-replace
2850    /// behavior, serialized so parallel first-use of the same extension family
2851    /// cannot publish over another thread's base snapshot.
2852    ///
2853    /// # Errors
2854    ///
2855    /// Returns [`tenferro_runtime::Error::RuntimeState`] when runtime
2856    /// reconfiguration fails or the extension module transaction is invalid.
2857    pub fn install_extension_module(
2858        &self,
2859        module: Arc<dyn ExtensionModule>,
2860    ) -> Result<RuntimeEpoch> {
2861        let snapshot = self.runtime.snapshot().map_err(|source| {
2862            runtime_state_source("EagerRuntime::install_extension_module", source)
2863        })?;
2864        if snapshot.has_extension_module_identical(&module) {
2865            return Ok(snapshot.epoch());
2866        }
2867        let _install_guard = self.lock_extension_install()?;
2868        self.runtime
2869            .reconfigure(|edit| {
2870                edit.replace_extension_module(module)?;
2871                Ok(())
2872            })
2873            .map_err(|source| {
2874                runtime_state_source("EagerRuntime::install_extension_module", source)
2875            })
2876    }
2877
2878    pub(crate) fn ensure_extension_module_for_engine(
2879        &self,
2880        module: Arc<dyn ExtensionModule>,
2881        family_id: &'static str,
2882        engine_id: &EngineId,
2883    ) -> Result<RuntimeEpoch> {
2884        let snapshot = self.runtime.snapshot().map_err(|source| {
2885            runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
2886        })?;
2887        if snapshot.has_extension_module_engine(module.module_id(), family_id, engine_id) {
2888            return Ok(snapshot.epoch());
2889        }
2890        let _install_guard = self.lock_extension_install()?;
2891        self.runtime
2892            .reconfigure(|edit| {
2893                edit.ensure_extension_module_for_engine(module, family_id, engine_id)?;
2894                Ok(())
2895            })
2896            .map_err(|source| {
2897                runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
2898            })
2899    }
2900
2901    pub(crate) fn runtime(&self) -> &Runtime {
2902        &self.runtime
2903    }
2904
2905    /// The extension target for `engine`, registering its kind first.
2906    pub(crate) fn eager_extension_target(
2907        &self,
2908        engine: &mut dyn ExecutionEngine,
2909    ) -> Result<EagerExtensionTarget> {
2910        let target = self.ensure_engine(engine)?;
2911        let backend_kind = eager_extension_backend_kind(&target).ok_or_else(|| {
2912            Error::unsupported(
2913                "EagerRuntime::eager_extension_target",
2914                ErrorPhase::Execution,
2915                "this engine kind has no eager extension target",
2916            )
2917        })?;
2918        let target = EagerExtensionTarget {
2919            engine_id: target.engine_id().clone(),
2920            backend_kind,
2921        };
2922        validate_eager_extension_target(&self.runtime, &target)?;
2923        Ok(target)
2924    }
2925
2926    /// Clear every cache owned by this eager context.
2927    ///
2928    /// # Examples
2929    ///
2930    /// ```
2931    /// use tenferro_ad::{EagerRuntime};
2932    ///
2933    /// let ctx = EagerRuntime::new()?;
2934    /// ctx.clear_caches()?;
2935    /// assert_eq!(ctx.cache_stats()?.ad_transforms.entries, 0);
2936    /// assert_eq!(ctx.cache_stats()?.prepared_derivatives.entries, 0);
2937    /// # Ok::<(), tenferro_ad::Error>(())
2938    /// ```
2939    ///
2940    /// # Errors
2941    ///
2942    /// Returns [`tenferro_runtime::Error::RuntimeState`] when either the
2943    /// extension cache or AD-transform cache is poisoned.
2944    pub fn clear_caches(&self) -> Result<()> {
2945        self.clear_ad_transform_caches()?;
2946        self.clear_prepared_derivative_cache()?;
2947        Ok(())
2948    }
2949
2950    /// Clear prepared derivative program cache entries.
2951    ///
2952    /// # Examples
2953    ///
2954    /// ```rust
2955    /// use tenferro_ad::EagerRuntime;
2956    ///
2957    /// let ctx = EagerRuntime::new()?;
2958    /// ctx.clear_prepared_derivative_cache()?;
2959    /// assert_eq!(ctx.cache_stats()?.prepared_derivatives.entries, 0);
2960    /// # Ok::<(), tenferro_ad::Error>(())
2961    /// ```
2962    ///
2963    /// # Errors
2964    ///
2965    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
2966    /// derivative cache lock is poisoned.
2967    pub fn clear_prepared_derivative_cache(&self) -> Result<()> {
2968        self.lock_prepared_derivative_cache()?.clear();
2969        Ok(())
2970    }
2971
2972    /// Return eager runtime cache-entry and retained-byte stats.
2973    ///
2974    /// # Examples
2975    ///
2976    /// ```
2977    /// use tenferro_ad::{EagerRuntime};
2978    ///
2979    /// let ctx = EagerRuntime::new()?;
2980    /// let stats = ctx.cache_stats()?;
2981    /// assert_eq!(stats.ad_transforms.entries, 0);
2982    /// assert_eq!(stats.prepared_derivatives.entries, 0);
2983    /// # Ok::<(), tenferro_ad::Error>(())
2984    /// ```
2985    ///
2986    /// # Errors
2987    ///
2988    /// Returns [`tenferro_runtime::Error::RuntimeState`] when a cache or
2989    /// AD-transform cache lock is poisoned.
2990    pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
2991        Ok(EagerRuntimeCacheStats {
2992            ad_transforms: self.ad_transform_cache.stats()?,
2993            prepared_derivatives: self.lock_prepared_derivative_cache()?.stats(),
2994        })
2995    }
2996
2997    /// Return the AD transform cache retention limits.
2998    ///
2999    /// # Examples
3000    ///
3001    /// ```
3002    /// use tenferro_ad::EagerRuntime;
3003    ///
3004    /// let ctx = EagerRuntime::new()?;
3005    /// assert!(ctx.ad_transform_cache_limits()?.max_entries().get() > 0);
3006    /// # Ok::<(), tenferro_ad::Error>(())
3007    /// ```
3008    ///
3009    /// # Errors
3010    ///
3011    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
3012    /// cache lock is poisoned.
3013    pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
3014        self.ad_transform_cache.limits()
3015    }
3016
3017    /// Replace AD transform cache retention limits.
3018    ///
3019    /// # Examples
3020    ///
3021    /// ```
3022    /// use std::num::NonZeroUsize;
3023    /// use tenferro_ad::{AdTransformCacheLimits, EagerRuntime};
3024    ///
3025    /// let ctx = EagerRuntime::new()?;
3026    /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
3027    /// ctx.set_ad_transform_cache_limits(limits)?;
3028    /// assert_eq!(ctx.ad_transform_cache_limits()?, limits);
3029    /// # Ok::<(), tenferro_ad::Error>(())
3030    /// ```
3031    ///
3032    /// # Errors
3033    ///
3034    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
3035    /// cache lock is poisoned while updating limits.
3036    pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
3037        self.ad_transform_cache.set_limits(limits)
3038    }
3039
3040    /// Clear AD transform cache entries visible through this eager runtime.
3041    ///
3042    /// # Examples
3043    ///
3044    /// ```
3045    /// use tenferro_ad::EagerRuntime;
3046    ///
3047    /// let ctx = EagerRuntime::new()?;
3048    /// ctx.clear_ad_transform_caches()?;
3049    /// assert_eq!(ctx.cache_stats()?.ad_transforms.entries, 0);
3050    /// # Ok::<(), tenferro_ad::Error>(())
3051    /// ```
3052    ///
3053    /// # Errors
3054    ///
3055    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
3056    /// cache lock is poisoned while clearing entries.
3057    pub fn clear_ad_transform_caches(&self) -> Result<()> {
3058        self.ad_transform_cache.clear()
3059    }
3060
3061    /// Return prepared derivative cache retention limits.
3062    ///
3063    /// # Examples
3064    ///
3065    /// ```rust
3066    /// use tenferro_ad::EagerRuntime;
3067    ///
3068    /// let ctx = EagerRuntime::new()?;
3069    /// assert!(ctx.prepared_derivative_cache_limits()?.max_entries().get() > 0);
3070    /// # Ok::<(), tenferro_ad::Error>(())
3071    /// ```
3072    ///
3073    /// # Errors
3074    ///
3075    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
3076    /// derivative cache lock is poisoned.
3077    pub fn prepared_derivative_cache_limits(&self) -> Result<AdTransformCacheLimits> {
3078        Ok(self.lock_prepared_derivative_cache()?.limits())
3079    }
3080
3081    /// Replace prepared derivative cache retention limits.
3082    ///
3083    /// # Examples
3084    ///
3085    /// ```rust
3086    /// use std::num::NonZeroUsize;
3087    /// use tenferro_ad::{AdTransformCacheLimits, EagerRuntime};
3088    ///
3089    /// let ctx = EagerRuntime::new()?;
3090    /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
3091    /// ctx.set_prepared_derivative_cache_limits(limits)?;
3092    /// assert_eq!(ctx.prepared_derivative_cache_limits()?, limits);
3093    /// # Ok::<(), tenferro_ad::Error>(())
3094    /// ```
3095    ///
3096    /// # Errors
3097    ///
3098    /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
3099    /// derivative cache lock is poisoned.
3100    pub fn set_prepared_derivative_cache_limits(
3101        &self,
3102        limits: AdTransformCacheLimits,
3103    ) -> Result<()> {
3104        self.lock_prepared_derivative_cache()?.set_limits(limits);
3105        Ok(())
3106    }
3107
3108    /// Run eager operations on `engine`.
3109    ///
3110    /// The callback receives an [`EagerSession`] that pairs this runtime with
3111    /// the borrowed engine; operations recorded through it are differentiable.
3112    /// Nothing is entered or locked: the engine is simply borrowed for the
3113    /// callback.
3114    ///
3115    /// # Examples
3116    ///
3117    /// ```
3118    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3119    /// use tenferro_cpu::CpuEngine;
3120    ///
3121    /// let ctx = EagerRuntime::new()?;
3122    /// let mut engine = CpuEngine::serial();
3123    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?, ctx.clone())?;
3124    /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?, ctx.clone())?;
3125    /// let z = ctx.with_eager_session(&mut engine, |s| s.add(&x, &y))?;
3126    /// assert_eq!(z.value()?.as_slice::<f64>()?, &[3.0]);
3127    /// # Ok::<(), tenferro_ad::Error>(())
3128    /// ```
3129    ///
3130    /// # Errors
3131    ///
3132    /// Returns [`Error::Unsupported`] when the engine kind cannot be registered
3133    /// with this runtime or its context identity conflicts with the existing
3134    /// registration, [`Error::RuntimeStateSource`] when the runtime snapshot or
3135    /// engine registration state is unavailable, and the callback's own error
3136    /// unchanged.
3137    pub fn with_eager_session<T, E: From<Error>>(
3138        self: &Arc<Self>,
3139        engine: &mut dyn ExecutionEngine,
3140        f: impl FnOnce(&mut EagerSession<'_>) -> std::result::Result<T, E>,
3141    ) -> std::result::Result<T, E> {
3142        self.ensure_engine(engine).map_err(E::from)?;
3143        f(&mut EagerSession {
3144            runtime: self,
3145            engine,
3146        })
3147    }
3148
3149    /// Run a compiled program on `engine` through this runtime.
3150    pub(crate) fn run_program(
3151        &self,
3152        engine: &mut dyn ExecutionEngine,
3153        program: &CompiledGraph,
3154        inputs: &[TensorRead<'_>],
3155    ) -> Result<Vec<Tensor>> {
3156        let target = self.ensure_engine(engine)?;
3157        let prepared = self
3158            .runtime
3159            .prepare(program, inputs, std::slice::from_ref(&target))?;
3160        run_prepared_on(&self.runtime, engine, &prepared, inputs)
3161    }
3162
3163    pub(crate) fn exec_extension_outputs_read(
3164        &self,
3165        engine: &mut dyn ExecutionEngine,
3166        op: &Arc<dyn tenferro_ops::ext_op::ExtensionOp>,
3167        inputs: &[TensorRead<'_>],
3168    ) -> Result<Vec<Tensor>> {
3169        self.ensure_engine(engine)?;
3170        profile_eager_op_section("exec_extension_outputs_read.exec_op", || {
3171            exec_extension_op_on_tensor_reads(op, inputs, engine, &self.runtime)
3172        })
3173    }
3174
3175    #[cfg(test)]
3176    pub(crate) fn exec_standard_graph_outputs(
3177        &self,
3178        engine: &mut dyn ExecutionEngine,
3179        graph: &Graph<StdTensorOp>,
3180        initial_data: HashMap<ValueKey<StdTensorOp>, Tensor>,
3181    ) -> Result<EagerGraphExecution> {
3182        let mut all_values = initial_data;
3183
3184        profile_eager_op_section("exec_graph.run", || {
3185            let exec: &mut dyn Engine = engine;
3186            (|| -> Result<()> {
3187                for op_node in graph.operations() {
3188                    let outputs = {
3189                        let input_values = op_node
3190                            .inputs
3191                            .iter()
3192                            .map(|input| {
3193                                let key = match input {
3194                                    ValueRef::Local(local_id) => &graph.values()[*local_id].key,
3195                                    ValueRef::External(key) => key,
3196                                };
3197                                all_values.get(key).ok_or_else(|| {
3198                                    Error::Internal(format!(
3199                                        "standard graph eager execution missing value for {key:?}"
3200                                    ))
3201                                })
3202                            })
3203                            .collect::<Result<Vec<_>>>()?;
3204                        let input_reads = input_values
3205                            .iter()
3206                            .map(|value| TensorRead::from_tensor(value))
3207                            .collect::<Vec<_>>();
3208                        exec_standard_op_on_tensor_reads_in_session(
3209                            &op_node.operation,
3210                            &input_reads,
3211                            exec,
3212                        )?
3213                    };
3214
3215                    if outputs.len() != op_node.outputs.len() {
3216                        return Err(Error::Internal(format!(
3217                            "standard graph eager execution expected {} outputs for {:?}, got {}",
3218                            op_node.outputs.len(),
3219                            op_node.operation,
3220                            outputs.len()
3221                        )));
3222                    }
3223
3224                    for (output_id, output) in op_node.outputs.iter().zip(outputs) {
3225                        let key = graph.values()[*output_id].key.clone();
3226                        all_values.insert(key, output);
3227                    }
3228                }
3229                Ok(())
3230            })()
3231        })?;
3232
3233        let outputs = graph
3234            .outputs()
3235            .iter()
3236            .map(|&output_id| {
3237                let key = &graph.values()[output_id].key;
3238                all_values
3239                    .get(key)
3240                    .ok_or_else(|| {
3241                        Error::Internal(format!(
3242                            "standard graph eager execution missing graph output {key:?}"
3243                        ))
3244                    })?
3245                    .duplicate()
3246                    .map_err(Error::from)
3247            })
3248            .collect::<Result<Vec<_>>>()?;
3249
3250        Ok(EagerGraphExecution { outputs })
3251    }
3252
3253    pub(crate) fn try_register_grad_slot(
3254        &self,
3255        key: &ValueKey<StdTensorOp>,
3256        slot: &GradSlot,
3257    ) -> Result<()> {
3258        insert_pruning_dead(
3259            &mut *self.lock_grad_slots()?,
3260            key.clone(),
3261            Arc::downgrade(slot),
3262        );
3263        Ok(())
3264    }
3265
3266    pub(crate) fn try_register_value_record(
3267        &self,
3268        key: &ValueKey<StdTensorOp>,
3269        record: &Arc<EagerTensorRecord>,
3270    ) -> Result<()> {
3271        insert_pruning_dead(
3272            &mut *self.lock_value_records()?,
3273            key.clone(),
3274            Arc::downgrade(record),
3275        );
3276        Ok(())
3277    }
3278
3279    pub(crate) fn value_record(
3280        &self,
3281        key: &ValueKey<StdTensorOp>,
3282    ) -> Result<Option<Arc<EagerTensorRecord>>> {
3283        let mut records = self.lock_value_records()?;
3284        let Some(record) = records.get(key).cloned() else {
3285            return Ok(None);
3286        };
3287        match record.upgrade() {
3288            Some(record) => Ok(Some(record)),
3289            None => {
3290                records.remove(key);
3291                Ok(None)
3292            }
3293        }
3294    }
3295
3296    /// Clear all live gradient slots tracked by this context.
3297    ///
3298    /// This resets the stored gradients to `None` without unregistering the
3299    /// tensors, so future `backward()` calls can accumulate again.
3300    ///
3301    /// # Examples
3302    ///
3303    /// ```
3304    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3305    ///
3306    /// let ctx = EagerRuntime::new()?;
3307    /// 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();
3308    /// 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();
3309    /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
3310    ///     let product = s.mul(&x, &y)?;
3311    ///     s.reduce_sum(&product, Some(&[0]))
3312    /// })?;
3313    /// let _ = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
3314    ///
3315    /// ctx.clear_grads()?;
3316    ///
3317    /// assert!(x.grad()?.is_none());
3318    /// assert!(y.grad()?.is_none());
3319    /// # Ok::<(), tenferro_ad::Error>(())
3320    /// ```
3321    ///
3322    /// # Errors
3323    ///
3324    /// Returns [`tenferro_runtime::Error::RuntimeState`] if a gradient-slot
3325    /// lock is poisoned while clearing live gradients.
3326    pub fn clear_grads(&self) -> Result<()> {
3327        let live_slots = {
3328            let mut live_slots = Vec::new();
3329            self.lock_grad_slots()?.retain(|_, slot| {
3330                if let Some(slot) = slot.upgrade() {
3331                    live_slots.push(slot);
3332                    true
3333                } else {
3334                    false
3335                }
3336            });
3337            live_slots
3338        };
3339
3340        let mut poisoned_slot = false;
3341        for slot in live_slots {
3342            match slot.lock() {
3343                Ok(mut current) => {
3344                    *current = None;
3345                }
3346                Err(_) => {
3347                    poisoned_slot = true;
3348                }
3349            }
3350        }
3351        if poisoned_slot {
3352            return Err(Error::runtime_state(
3353                "eager_gradient_slot",
3354                ErrorPhase::Execution,
3355                "lock poisoned",
3356            ));
3357        }
3358        Ok(())
3359    }
3360
3361    /// Import a concrete tensor into this context as an untracked constant.
3362    ///
3363    /// The returned tensor does not participate in gradient tracking.
3364    /// Use this for fixed masks, quadrature weights, physical constants,
3365    /// and other data that should not receive gradients. Like
3366    /// [`EagerSession::constant_from`], it performs no host/device transfer;
3367    /// use [`EagerSession::constant_from_host`] to upload host data into a
3368    /// device runtime.
3369    ///
3370    /// # Examples
3371    ///
3372    /// ```
3373    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3374    ///
3375    /// let ctx = EagerRuntime::new()?;
3376    /// let c = ctx.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
3377    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone())?;
3378    /// let z = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.add(&x, &c))?;
3379    ///
3380    /// assert_eq!(z.value()?.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
3381    /// # Ok::<(), tenferro_ad::Error>(())
3382    /// ```
3383    ///
3384    /// # Errors
3385    ///
3386    /// Returns [`tenferro_runtime::Error::RuntimeState`] when metadata cannot
3387    /// be registered or the backend lock is poisoned.
3388    pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
3389        EagerTensor::new_leaf(Arc::clone(self), tensor, false)
3390    }
3391
3392    /// Import a concrete tensor into this context as a trainable variable.
3393    ///
3394    /// The returned tensor participates in gradient tracking; its gradient
3395    /// slot is registered in this context.
3396    ///
3397    /// # Examples
3398    ///
3399    /// ```
3400    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3401    ///
3402    /// let ctx = EagerRuntime::new()?;
3403    /// let p = ctx.variable_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
3404    /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
3405    ///     let y = s.exp(&p)?;
3406    ///     s.reduce_sum(&y, Some(&[0]))
3407    /// })?;
3408    /// let _ = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
3409    ///
3410    /// let grad = p.grad().unwrap().unwrap();
3411    /// assert_eq!(grad.shape(), &[2]);
3412    /// # Ok::<(), tenferro_ad::Error>(())
3413    /// ```
3414    ///
3415    /// # Errors
3416    ///
3417    /// Returns [`tenferro_runtime::Error::RuntimeState`] when gradient metadata
3418    /// or the eager backend state cannot be registered.
3419    pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
3420        EagerTensor::new_leaf(Arc::clone(self), tensor, true)
3421    }
3422
3423    /// Gradient of a scalar eager output with respect to an eager tensor.
3424    ///
3425    /// Functional eager gradients return ordinary eager tensors and do not
3426    /// write into `grad()` slots. The returned tensor keeps a trace when the
3427    /// derivative computation depends on tracked eager values.
3428    ///
3429    /// # Examples
3430    ///
3431    /// ```
3432    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3433    ///
3434    /// let ctx = EagerRuntime::new()?;
3435    /// let x = EagerTensor::requires_grad_in(
3436    ///     Tensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap(),
3437    ///     ctx.clone(),
3438    /// )?;
3439    /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
3440    /// let dx = ctx.grad(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x)?;
3441    /// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[6.0]);
3442    /// # Ok::<(), tenferro_ad::Error>(())
3443    /// ```
3444    ///
3445    /// # Errors
3446    ///
3447    /// Returns [`tenferro_runtime::Error::NonScalarGrad`] for a non-scalar
3448    /// output, [`Error::ContextMismatch`] for tensors from another runtime,
3449    /// [`Error::UnsupportedAdRule`] when an AD rule is unavailable, or a typed
3450    /// validation/backend error from eager execution. An inactive `wrt` returns
3451    /// [`Error::Validation`] with `argument: "wrt"`; use
3452    /// [`grad_optional`](Self::grad_optional) to observe that state.
3453    pub fn grad(
3454        self: &Arc<Self>,
3455        engine: &mut dyn ExecutionEngine,
3456        output: &EagerTensor,
3457        wrt: &EagerTensor,
3458    ) -> Result<EagerTensor> {
3459        self.grad_optional(engine, output, wrt)?
3460            .ok_or_else(|| crate::traced::inactive_wrt_error("grad", &wrt.key))
3461    }
3462
3463    /// Gradient that returns `None` when `wrt` is inactive.
3464    ///
3465    /// # Examples
3466    ///
3467    /// ```
3468    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3469    ///
3470    /// let ctx = EagerRuntime::new()?;
3471    /// let x = EagerTensor::requires_grad_in(
3472    ///     Tensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap(),
3473    ///     ctx.clone(),
3474    /// )?;
3475    /// let y = EagerTensor::requires_grad_in(
3476    ///     Tensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap(),
3477    ///     ctx.clone(),
3478    /// )?;
3479    /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&y, &y))?;
3480    /// assert!(ctx.grad_optional(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x)?.is_none());
3481    /// # Ok::<(), tenferro_ad::Error>(())
3482    /// ```
3483    ///
3484    /// # Errors
3485    ///
3486    /// Returns [`tenferro_runtime::Error::NonScalarGrad`] for a non-scalar
3487    /// output, [`Error::ContextMismatch`] for a foreign runtime, or a typed
3488    /// validation/backend/runtime-state error from eager execution.
3489    pub fn grad_optional(
3490        self: &Arc<Self>,
3491        engine: &mut dyn ExecutionEngine,
3492        output: &EagerTensor,
3493        wrt: &EagerTensor,
3494    ) -> Result<Option<EagerTensor>> {
3495        if !output.shape().is_empty() {
3496            return Err(Error::NonScalarGrad {
3497                shape: output.shape().to_vec(),
3498            });
3499        }
3500
3501        let value = output.to_tensor(&mut *engine)?;
3502        let seed = one_like_tensor(&value, &mut *engine)?;
3503        let seed = EagerTensor::new_result(Arc::clone(self), eager_val_key(), seed, false, None)?;
3504        self.vjp_optional(engine, output, wrt, &seed)
3505    }
3506
3507    /// Reverse-mode vector-Jacobian product for eager tensors.
3508    ///
3509    /// # Examples
3510    ///
3511    /// ```
3512    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3513    ///
3514    /// let ctx = EagerRuntime::new()?;
3515    /// let x = EagerTensor::requires_grad_in(
3516    ///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
3517    ///     ctx.clone(),
3518    /// )?;
3519    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
3520    /// let seed = EagerTensor::from_tensor_in(
3521    ///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 1.0]).unwrap(),
3522    ///     ctx.clone(),
3523    /// )?;
3524    /// let dx = ctx.vjp(&mut tenferro_cpu::CpuEngine::serial(), &y, &x, &seed)?;
3525    /// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
3526    /// # Ok::<(), tenferro_ad::Error>(())
3527    /// ```
3528    ///
3529    /// # Errors
3530    ///
3531    /// Returns [`Error::ContextMismatch`] for tensors from different eager
3532    /// runtimes, [`Error::Validation`] when the cotangent shape or dtype does
3533    /// not match the output, [`Error::UnsupportedAdRule`] when a rule is not
3534    /// registered, or a typed backend/runtime-state error. An inactive `wrt`
3535    /// returns [`Error::Validation`] with `argument: "wrt"`; use
3536    /// [`vjp_optional`](Self::vjp_optional) to observe that state.
3537    pub fn vjp(
3538        self: &Arc<Self>,
3539        engine: &mut dyn ExecutionEngine,
3540        output: &EagerTensor,
3541        wrt: &EagerTensor,
3542        cotangent: &EagerTensor,
3543    ) -> Result<EagerTensor> {
3544        self.vjp_optional(engine, output, wrt, cotangent)?
3545            .ok_or_else(|| crate::traced::inactive_wrt_error("vjp", &wrt.key))
3546    }
3547
3548    /// Reverse-mode vector-Jacobian product that returns `None` for inactive inputs.
3549    ///
3550    /// # Examples
3551    ///
3552    /// ```
3553    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3554    ///
3555    /// let ctx = EagerRuntime::new()?;
3556    /// let x = EagerTensor::requires_grad_in(
3557    ///     Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
3558    ///     ctx.clone(),
3559    /// )?;
3560    /// let y = EagerTensor::requires_grad_in(
3561    ///     Tensor::from_vec_col_major(vec![1], vec![4.0_f64]).unwrap(),
3562    ///     ctx.clone(),
3563    /// )?;
3564    /// let seed = EagerTensor::from_tensor_in(
3565    ///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
3566    ///     ctx.clone(),
3567    /// )?;
3568    /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&y, &y))?;
3569    /// assert!(ctx.vjp_optional(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x, &seed)?.is_none());
3570    /// # Ok::<(), tenferro_ad::Error>(())
3571    /// ```
3572    ///
3573    /// # Errors
3574    ///
3575    /// Returns [`Error::ContextMismatch`] for tensors from different eager
3576    /// runtimes, [`Error::Validation`] when the cotangent shape or dtype does
3577    /// not match the output, [`Error::UnsupportedAdRule`] when a rule is not
3578    /// registered, or a typed backend/runtime-state error.
3579    pub fn vjp_optional(
3580        self: &Arc<Self>,
3581        engine: &mut dyn ExecutionEngine,
3582        output: &EagerTensor,
3583        wrt: &EagerTensor,
3584        cotangent: &EagerTensor,
3585    ) -> Result<Option<EagerTensor>> {
3586        validate_same_runtime(self, output, "vjp output")?;
3587        validate_same_runtime(self, wrt, "vjp wrt")?;
3588        validate_same_runtime(self, cotangent, "vjp cotangent")?;
3589        validate_seed_tensor("vjp", output, cotangent)?;
3590        Ok(
3591            semantic_eager_vjp_many(self, engine, output, &[wrt], cotangent)?
3592                .pop()
3593                .flatten(),
3594        )
3595    }
3596
3597    /// Forward-mode Jacobian-vector product for eager tensors.
3598    ///
3599    /// # Examples
3600    ///
3601    /// ```
3602    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3603    ///
3604    /// let ctx = EagerRuntime::new()?;
3605    /// let x = EagerTensor::requires_grad_in(
3606    ///     Tensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap(),
3607    ///     ctx.clone(),
3608    /// )?;
3609    /// let tangent = EagerTensor::from_tensor_in(
3610    ///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
3611    ///     ctx.clone(),
3612    /// )?;
3613    /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
3614    /// let dy = ctx.jvp(&mut tenferro_cpu::CpuEngine::serial(), &y, &x, &tangent)?;
3615    /// assert_eq!(dy.value()?.as_slice::<f64>().unwrap(), &[6.0]);
3616    /// # Ok::<(), tenferro_ad::Error>(())
3617    /// ```
3618    ///
3619    /// # Errors
3620    ///
3621    /// Returns [`Error::ContextMismatch`] for tensors from different eager
3622    /// runtimes, [`Error::Validation`] when the tangent shape or dtype does not
3623    /// match `wrt`, [`Error::UnsupportedAdRule`] when a rule is unavailable, or
3624    /// a typed backend/runtime-state error. An inactive `wrt` returns
3625    /// [`Error::Validation`] with `argument: "wrt"`; use
3626    /// [`jvp_optional`](Self::jvp_optional) to observe that state.
3627    pub fn jvp(
3628        self: &Arc<Self>,
3629        engine: &mut dyn ExecutionEngine,
3630        output: &EagerTensor,
3631        wrt: &EagerTensor,
3632        tangent: &EagerTensor,
3633    ) -> Result<EagerTensor> {
3634        self.jvp_optional(engine, output, wrt, tangent)?
3635            .ok_or_else(|| crate::traced::inactive_wrt_error("jvp", &wrt.key))
3636    }
3637
3638    /// Forward-mode Jacobian-vector product that returns `None` for inactive outputs.
3639    ///
3640    /// # Examples
3641    ///
3642    /// ```
3643    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3644    ///
3645    /// let ctx = EagerRuntime::new()?;
3646    /// let x = EagerTensor::requires_grad_in(
3647    ///     Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
3648    ///     ctx.clone(),
3649    /// )?;
3650    /// let y = EagerTensor::requires_grad_in(
3651    ///     Tensor::from_vec_col_major(vec![1], vec![4.0_f64]).unwrap(),
3652    ///     ctx.clone(),
3653    /// )?;
3654    /// let tangent = EagerTensor::from_tensor_in(
3655    ///     Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
3656    ///     ctx.clone(),
3657    /// )?;
3658    /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&y, &y))?;
3659    /// assert!(ctx.jvp_optional(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x, &tangent)?.is_none());
3660    /// # Ok::<(), tenferro_ad::Error>(())
3661    /// ```
3662    ///
3663    /// # Errors
3664    ///
3665    /// Returns [`Error::ContextMismatch`] for tensors from different eager
3666    /// runtimes, [`Error::Validation`] when the tangent shape or dtype does not
3667    /// match `wrt`, [`Error::UnsupportedAdRule`] when a rule is unavailable, or
3668    /// a typed backend/runtime-state error.
3669    pub fn jvp_optional(
3670        self: &Arc<Self>,
3671        engine: &mut dyn ExecutionEngine,
3672        output: &EagerTensor,
3673        wrt: &EagerTensor,
3674        tangent: &EagerTensor,
3675    ) -> Result<Option<EagerTensor>> {
3676        validate_same_runtime(self, output, "jvp output")?;
3677        validate_same_runtime(self, wrt, "jvp wrt")?;
3678        validate_same_runtime(self, tangent, "jvp tangent")?;
3679        validate_seed_tensor("jvp", wrt, tangent)?;
3680        // Unification 7: semantic path is the only JVP path.
3681        match semantic_eager_jvp_optional(self, engine, output, wrt, tangent)? {
3682            Some(result) => Ok(result),
3683            None => Ok(None),
3684        }
3685    }
3686
3687    fn store_grads(
3688        &self,
3689        cotangents: &HashMap<ValueKey<StdTensorOp>, Tensor>,
3690        session: &mut dyn Engine,
3691    ) -> Result<()> {
3692        let mut updates = Vec::new();
3693
3694        {
3695            let mut slots = self.lock_grad_slots()?;
3696            slots.retain(|key, slot| {
3697                let Some(slot) = slot.upgrade() else {
3698                    return false;
3699                };
3700
3701                if let Some(incoming) = cotangents.get(key) {
3702                    updates.push((slot, incoming));
3703                }
3704
3705                true
3706            });
3707        }
3708
3709        for (slot, incoming) in updates {
3710            let mut current = slot.lock().map_err(|_| {
3711                Error::runtime_state(
3712                    "eager_gradient_slot",
3713                    ErrorPhase::Execution,
3714                    "lock poisoned",
3715                )
3716            })?;
3717            let next = match current.as_ref() {
3718                Some(existing) => {
3719                    let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
3720                    let incoming_read = TensorRead::from_tensor(incoming);
3721                    let tensor = session
3722                        .add_read(existing_read, incoming_read)
3723                        .map_err(Error::from)?;
3724                    AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
3725                }
3726                None => {
3727                    let duplicate = session
3728                        .to_contiguous_read(TensorRead::from_tensor(incoming))
3729                        .map_err(Error::from)?;
3730                    AdValueRecord::from_tensor(duplicate, "EagerRuntime::store_grads")?
3731                }
3732            };
3733            *current = Some(next);
3734        }
3735
3736        Ok(())
3737    }
3738}
3739
3740#[derive(Clone, Debug, PartialEq, Eq, Hash)]
3741struct PreparedDerivativeCacheKey {
3742    /// Engines with the same target share prepared derivative programs.
3743    target: EngineTarget,
3744    semantic_fingerprint: SemanticFingerprint,
3745    runtime_epoch: RuntimeEpoch,
3746    active_inputs: Box<[bool]>,
3747    input_metadata: Box<[ProgramValueMetadata]>,
3748}
3749
3750/// Cached prepared derivative: program + index metadata.
3751#[derive(Debug)]
3752struct PreparedDerivative {
3753    program: Arc<CompiledGraph>,
3754    execution_program: Arc<CompiledGraph>,
3755    saved_input_indices: Vec<usize>,
3756    prepared: Arc<Prepared>,
3757    seed_input_index: usize,
3758    derivative_output_indices: Box<[Option<usize>]>,
3759}
3760
3761#[derive(Debug)]
3762struct PreparedDerivativeCache {
3763    limits: AdTransformCacheLimits,
3764    entries: LruCache<PreparedDerivativeCacheKey, PreparedDerivativeCacheEntry>,
3765    stats: CacheStats,
3766}
3767
3768impl PreparedDerivativeCache {
3769    fn limits(&self) -> AdTransformCacheLimits {
3770        self.limits
3771    }
3772
3773    fn set_limits(&mut self, limits: AdTransformCacheLimits) {
3774        self.limits = limits;
3775        self.evict_to_limits();
3776    }
3777
3778    fn clear(&mut self) {
3779        let clears = self.stats.clears.saturating_add(1);
3780        self.entries.clear();
3781        self.stats = CacheStats {
3782            clears,
3783            ..CacheStats::empty()
3784        };
3785    }
3786
3787    fn stats(&self) -> CacheStats {
3788        self.stats
3789    }
3790
3791    fn get(&mut self, key: &PreparedDerivativeCacheKey) -> Option<Arc<PreparedDerivative>> {
3792        match self.entries.get(key) {
3793            Some(entry) => {
3794                self.stats.hits = self.stats.hits.saturating_add(1);
3795                Some(Arc::clone(&entry.value))
3796            }
3797            None => {
3798                self.stats.misses = self.stats.misses.saturating_add(1);
3799                None
3800            }
3801        }
3802    }
3803
3804    fn insert(&mut self, key: PreparedDerivativeCacheKey, value: Arc<PreparedDerivative>) {
3805        let retained_bytes = prepared_derivative_cache_entry_retained_bytes(&key, value.as_ref());
3806        let entry = PreparedDerivativeCacheEntry {
3807            value,
3808            retained_bytes,
3809        };
3810        self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
3811        if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
3812            self.stats.retained_bytes = self
3813                .stats
3814                .retained_bytes
3815                .saturating_sub(old_entry.retained_bytes);
3816        }
3817        self.stats.entries = self.entries.len();
3818        self.evict_to_limits();
3819    }
3820
3821    fn evict_to_limits(&mut self) {
3822        while self.entries.len() > self.limits.max_entries().get()
3823            || self
3824                .limits
3825                .max_retained_bytes()
3826                .is_some_and(|limit| self.stats.retained_bytes > limit.get())
3827        {
3828            let Some((_key, entry)) = self.entries.pop_lru() else {
3829                break;
3830            };
3831            self.stats.retained_bytes = self
3832                .stats
3833                .retained_bytes
3834                .saturating_sub(entry.retained_bytes);
3835            self.stats.evictions = self.stats.evictions.saturating_add(1);
3836        }
3837        self.stats.entries = self.entries.len();
3838    }
3839}
3840
3841impl Default for PreparedDerivativeCache {
3842    fn default() -> Self {
3843        Self {
3844            limits: AdTransformCacheLimits::default(),
3845            entries: LruCache::unbounded(),
3846            stats: CacheStats::empty(),
3847        }
3848    }
3849}
3850
3851#[derive(Debug)]
3852struct PreparedDerivativeCacheEntry {
3853    value: Arc<PreparedDerivative>,
3854    retained_bytes: usize,
3855}
3856
3857fn prepared_derivative_cache_entry_retained_bytes(
3858    key: &PreparedDerivativeCacheKey,
3859    value: &PreparedDerivative,
3860) -> usize {
3861    size_of::<PreparedDerivativeCacheKey>()
3862        .saturating_add(size_of_val(key.active_inputs.as_ref()))
3863        .saturating_add(size_of_val(value.derivative_output_indices.as_ref()))
3864        .saturating_add(size_of_val(value.saved_input_indices.as_slice()))
3865        .saturating_add(compiled_graph_retained_bytes(
3866            value.execution_program.as_ref(),
3867        ))
3868        .saturating_add(
3869            key.input_metadata
3870                .len()
3871                .saturating_mul(size_of::<ProgramValueMetadata>()),
3872        )
3873        .saturating_add(size_of::<PreparedDerivative>())
3874        .saturating_add(compiled_graph_retained_bytes(value.program.as_ref()))
3875        .saturating_add(prepared_compiled_graph_retained_bytes(
3876            value.prepared.as_ref(),
3877            value.program.as_ref(),
3878        ))
3879}
3880
3881fn prepared_compiled_graph_retained_bytes(
3882    prepared: &Prepared,
3883    derivative_program: &CompiledGraph,
3884) -> usize {
3885    size_of_val(prepared).saturating_add(compiled_graph_retained_bytes(derivative_program))
3886}
3887
3888fn compiled_graph_retained_bytes(program: &CompiledGraph) -> usize {
3889    size_of::<CompiledGraph>()
3890        .saturating_add(size_of_val(program.input_keys()))
3891        .saturating_add(program.bindings().len().saturating_mul(size_of::<usize>()))
3892        .saturating_add(semantic_program_retained_bytes(program.program()))
3893}
3894
3895fn semantic_program_retained_bytes(program: &SemanticProgram) -> usize {
3896    size_of::<SemanticProgram>()
3897        .saturating_add(size_of_val(program.inputs()))
3898        .saturating_add(size_of_val(program.outputs()))
3899        .saturating_add(
3900            program
3901                .operations()
3902                .len()
3903                .saturating_mul(size_of::<usize>()),
3904        )
3905        .saturating_add(
3906            program
3907                .shape_guards()
3908                .len()
3909                .saturating_mul(size_of::<usize>()),
3910        )
3911}
3912
3913fn semantic_eager_vjp_many(
3914    ctx: &Arc<EagerRuntime>,
3915    engine: &mut dyn ExecutionEngine,
3916    output: &EagerTensor,
3917    wrts: &[&EagerTensor],
3918    cotangent: &EagerTensor,
3919) -> Result<Vec<Option<EagerTensor>>> {
3920    if !eager_semantic_vjp_enabled() || wrts.is_empty() {
3921        return Ok(vec![None; wrts.len()]);
3922    }
3923    let Some(raw_output_trace) = output.semantic_trace.as_ref() else {
3924        return Ok(vec![None; wrts.len()]);
3925    };
3926    if !wrts.iter().any(|wrt| {
3927        wrt.semantic_trace
3928            .as_ref()
3929            .and_then(TracedTensor::input_key)
3930            .is_some_and(|key| raw_output_trace.has_attached_input_key(&key))
3931    }) {
3932        return Ok(vec![None; wrts.len()]);
3933    }
3934
3935    // First AD request on this output: run the deferred graph analysis over
3936    // the whole raw carrier chain once (metadata registration + constraint
3937    // scopes), so `compile_ad_source` sees the same analyzed graph the eager
3938    // forward used to append.
3939    let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
3940    let saved = output
3941        .trace
3942        .as_ref()
3943        .map(EagerTrace::collect)
3944        .unwrap_or_default();
3945    let mut source_outputs = vec![&output_trace];
3946    source_outputs.extend(saved.iter().map(|value| &value.trace));
3947
3948    // Residual roots keep the semantic producers available to differentiation;
3949    // only the execution-only derivative replaces their numerical values.
3950    let mut compiler = GraphCompiler::new();
3951    let source =
3952        tenferro_runtime::ad_support::compile_ad_source_many(&mut compiler, &source_outputs)?;
3953    if source.output_count() != source_outputs.len()
3954        || source.input_keys().len() != source.input_count()
3955        || source.bindings().len() != source.input_count()
3956    {
3957        return Ok(vec![None; wrts.len()]);
3958    }
3959    let wrt_input_indices = wrts
3960        .iter()
3961        .map(|wrt| {
3962            let key = wrt.semantic_trace.as_ref()?.input_key()?;
3963            source.input_key_index(&key)
3964        })
3965        .collect::<Vec<_>>();
3966    let mut active_inputs = vec![false; source.input_count()];
3967    for &index in wrt_input_indices.iter().flatten() {
3968        active_inputs[index] = true;
3969    }
3970    if !active_inputs.iter().any(|&active| active) {
3971        return Ok(vec![None; wrts.len()]);
3972    }
3973
3974    // One transform and execution for the complete active set shares primal
3975    // work and cotangents between leaves instead of replaying them per target.
3976    // S2: check prepared-derivative cache before AD transform + compile_frozen.
3977    let target = ctx.ensure_engine(engine)?;
3978    let cache_key = PreparedDerivativeCacheKey {
3979        target: target.clone(),
3980        semantic_fingerprint: source.program().semantic_fingerprint(),
3981        runtime_epoch: ctx.runtime.epoch().map_err(|source| {
3982            Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::Execution, source)
3983        })?,
3984        active_inputs: active_inputs.clone().into_boxed_slice(),
3985        input_metadata: source.frozen_program().input_metadata_with_bound_shapes(),
3986    };
3987    let prepared = { ctx.lock_prepared_derivative_cache()?.get(&cache_key) };
3988    let (
3989        seed_input_index,
3990        derivative_output_indices,
3991        derivative_program,
3992        execution_program,
3993        saved_input_indices,
3994        prepared_runtime,
3995    ) = if let Some(prepared) = prepared {
3996        (
3997            prepared.seed_input_index,
3998            prepared.derivative_output_indices.clone(),
3999            Arc::clone(&prepared.program),
4000            Arc::clone(&prepared.execution_program),
4001            prepared.saved_input_indices.clone(),
4002            Some(Arc::clone(&prepared.prepared)),
4003        )
4004    } else {
4005        let mut active_outputs = vec![false; source.output_count()];
4006        active_outputs[0] = true;
4007        let ad = AdContext::with_rules_and_transform_cache(
4008            ctx.semantic_extension_rules.clone(),
4009            Arc::clone(&ctx.ad_transform_cache),
4010        );
4011        let derivative = ad
4012            .vjp_program(source.frozen_program(), &active_inputs, &active_outputs)
4013            .map_err(|source| {
4014                Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::GraphBuild, source)
4015            })?;
4016        let seed_input_index = derivative
4017            .derivative_input_indices()
4018            .first()
4019            .copied()
4020            .flatten();
4021        let Some(seed_input_index) = seed_input_index else {
4022            return Ok(vec![None; wrts.len()]);
4023        };
4024        let derivative_output_indices = derivative
4025            .derivative_output_indices()
4026            .to_vec()
4027            .into_boxed_slice();
4028        let program = Arc::new(compiler.compile_frozen_program(derivative.frozen())?);
4029        let (execution_program, saved_input_indices) = if saved.is_empty() {
4030            (Arc::clone(&program), Vec::new())
4031        } else {
4032            let (execution, indices) = crate::semantic_transform::semantic_vjp_with_saved_outputs(
4033                source.frozen_program(),
4034                &active_inputs,
4035                &active_outputs,
4036                &ctx.semantic_extension_rules,
4037                &(1..source.output_count()).collect::<Vec<_>>(),
4038            )
4039            .map_err(|source| {
4040                Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::GraphBuild, source)
4041            })?;
4042            (
4043                Arc::new(compiler.compile_frozen_program(execution.frozen())?),
4044                indices,
4045            )
4046        };
4047        (
4048            seed_input_index,
4049            derivative_output_indices,
4050            program,
4051            execution_program,
4052            saved_input_indices,
4053            None,
4054        )
4055    };
4056
4057    let cotangent_tensor = Arc::new(RetainedValue::from_tensor(
4058        cotangent.to_tensor(&mut *engine)?,
4059    ));
4060    let input_count = execution_program.input_count();
4061    let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
4062    // Residuals, primal bindings and the seed are staged on the caller's engine.
4063    (|| -> Result<()> {
4064        let session: &mut dyn Engine = &mut *engine;
4065        for (value, &index) in saved.iter().zip(&saved_input_indices) {
4066            let read = value.value.tensor_read("eager residual")?;
4067            let Some(slot) = owned_inputs.get_mut(index) else {
4068                return Err(Error::Internal(format!(
4069                    "semantic eager VJP residual index {index} is outside {input_count} inputs"
4070                )));
4071            };
4072            *slot = Some(session.to_contiguous_read(read)?);
4073        }
4074        for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
4075            let Some(slot) = owned_inputs.get_mut(source_input_index) else {
4076                return Err(Error::Internal(format!(
4077                    "semantic eager VJP derivative program has no primal input slot {source_input_index}"
4078                )));
4079            };
4080            *slot = Some(copy_value_in_session(session, tensor)?);
4081        }
4082        let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
4083            return Err(Error::Internal(format!(
4084                "semantic eager VJP seed input index {seed_input_index} is outside {input_count} inputs"
4085            )));
4086        };
4087        *slot = Some(copy_value_in_session(session, cotangent_tensor.as_ref())?);
4088        Ok(())
4089    })()?;
4090    let input_refs = owned_inputs
4091        .iter()
4092        .enumerate()
4093        .map(|(index, tensor)| {
4094            tensor.as_ref().ok_or_else(|| {
4095                Error::Internal(format!(
4096                    "semantic eager VJP derivative input {index} was not populated"
4097                ))
4098            })
4099        })
4100        .collect::<Result<Vec<_>>>()?;
4101    let input_reads = input_refs
4102        .iter()
4103        .map(|tensor| TensorRead::from_tensor(tensor))
4104        .collect::<Vec<_>>();
4105    let prepared_runtime = if let Some(prepared_runtime) = prepared_runtime {
4106        prepared_runtime
4107    } else {
4108        let prepared_runtime = Arc::new(ctx.runtime.prepare(
4109            &execution_program,
4110            &input_reads,
4111            std::slice::from_ref(&target),
4112        )?);
4113        let entry = Arc::new(PreparedDerivative {
4114            program: Arc::clone(&derivative_program),
4115            execution_program: Arc::clone(&execution_program),
4116            saved_input_indices: saved_input_indices.clone(),
4117            prepared: Arc::clone(&prepared_runtime),
4118            seed_input_index,
4119            derivative_output_indices: derivative_output_indices.clone(),
4120        });
4121        ctx.lock_prepared_derivative_cache()?
4122            .insert(cache_key, entry);
4123        prepared_runtime
4124    };
4125    let mut outputs = run_prepared_on(&ctx.runtime, engine, &prepared_runtime, &input_reads)?
4126        .into_iter()
4127        .map(Some)
4128        .collect::<Vec<_>>();
4129    let cotangent_trace =
4130        TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&cotangent_tensor))?;
4131
4132    #[cfg(test)]
4133    EAGER_SEMANTIC_VJP_EXECUTIONS.fetch_add(1, Ordering::Relaxed);
4134
4135    wrts.iter()
4136        .zip(wrt_input_indices)
4137        .map(|(wrt, input_index)| {
4138            let Some(derivative_output_index) = input_index
4139                .and_then(|index| derivative_output_indices.get(index).copied().flatten())
4140            else {
4141                return Ok(None);
4142            };
4143            let result = outputs
4144                .get_mut(derivative_output_index)
4145                .and_then(Option::take)
4146                .ok_or_else(|| {
4147                    Error::Internal(format!(
4148                "semantic eager VJP derivative output {derivative_output_index} unavailable"
4149            ))
4150                })?;
4151            let wrt_trace = wrt.semantic_trace.as_ref().ok_or_else(|| {
4152                Error::Internal("active eager VJP input has no semantic trace".into())
4153            })?;
4154            let semantic_trace = derivative_trace_from_frozen_program(
4155                &source,
4156                derivative_program.frozen_program(),
4157                derivative_output_index,
4158                &[(seed_input_index, Arc::clone(&cotangent_tensor))],
4159                &[&output_trace, wrt_trace, &cotangent_trace],
4160                None,
4161                "semantic_eager_vjp",
4162            )?;
4163            Ok(Some(EagerTensor::new_result_with_semantic_trace(
4164                Arc::clone(ctx),
4165                eager_val_key(),
4166                result,
4167                true,
4168                output.trace.clone(),
4169                Some(semantic_trace),
4170            )?))
4171        })
4172        .collect()
4173}
4174
4175fn semantic_eager_jvp_optional(
4176    ctx: &Arc<EagerRuntime>,
4177    engine: &mut dyn ExecutionEngine,
4178    output: &EagerTensor,
4179    wrt: &EagerTensor,
4180    tangent: &EagerTensor,
4181) -> Result<Option<Option<EagerTensor>>> {
4182    if !eager_semantic_vjp_enabled() {
4183        return Ok(None);
4184    }
4185    let (Some(raw_output_trace), Some(wrt_trace)) =
4186        (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
4187    else {
4188        return Ok(None);
4189    };
4190    let Some(wrt_key) = wrt_trace.input_key() else {
4191        return Ok(None);
4192    };
4193    if !raw_output_trace.has_attached_input_key(&wrt_key) {
4194        return Ok(None);
4195    }
4196
4197    // First AD request on this output: run the deferred graph analysis once
4198    // over the whole raw carrier chain before compiling.
4199    let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
4200
4201    let mut compiler = GraphCompiler::new();
4202    let source = compile_ad_source(&mut compiler, &output_trace)?;
4203    if source.output_count() != 1
4204        || source.input_keys().len() != source.input_count()
4205        || source.bindings().len() != source.input_count()
4206    {
4207        return Ok(None);
4208    }
4209    let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
4210        return Ok(None);
4211    };
4212
4213    let mut active_inputs = vec![false; source.input_count()];
4214    if let Some(active) = active_inputs.get_mut(wrt_input_index) {
4215        *active = true;
4216    } else {
4217        return Ok(None);
4218    }
4219    let ad = AdContext::with_rules_and_transform_cache(
4220        ctx.semantic_extension_rules.clone(),
4221        Arc::clone(&ctx.ad_transform_cache),
4222    );
4223    let derivative = ad
4224        .jvp_program(source.frozen_program(), &active_inputs)
4225        .map_err(|source| {
4226            Error::runtime_state_source("semantic_eager_jvp", ErrorPhase::GraphBuild, source)
4227        })?;
4228    // derivative_input_indices maps source input → derivative seed input.
4229    let Some(seed_input_index) = derivative
4230        .derivative_input_indices()
4231        .get(wrt_input_index)
4232        .copied()
4233        .flatten()
4234    else {
4235        return Ok(Some(None));
4236    };
4237    // derivative_output_indices maps source output → derivative output.
4238    // There is always exactly one source output (guarded above).
4239    let Some(derivative_output_index) = derivative
4240        .derivative_output_indices()
4241        .first()
4242        .copied()
4243        .flatten()
4244    else {
4245        return Ok(Some(None));
4246    };
4247
4248    let derivative_program = compiler.compile_frozen_program(derivative.frozen())?;
4249    let tangent_tensor = Arc::new(RetainedValue::from_tensor(tangent.to_tensor(&mut *engine)?));
4250    let input_count = derivative_program.input_count();
4251    let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
4252    // Primal bindings and the seed are staged on the caller's engine.
4253    (|| -> Result<()> {
4254        let session: &mut dyn Engine = &mut *engine;
4255        for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
4256            let Some(slot) = owned_inputs.get_mut(source_input_index) else {
4257                return Err(Error::Internal(format!(
4258                    "semantic eager JVP derivative program has no primal input slot {source_input_index}"
4259                )));
4260            };
4261            *slot = Some(copy_value_in_session(session, tensor)?);
4262        }
4263        let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
4264            return Err(Error::Internal(format!(
4265                "semantic eager JVP seed input index {seed_input_index} is outside {input_count} inputs"
4266            )));
4267        };
4268        *slot = Some(copy_value_in_session(session, tangent_tensor.as_ref())?);
4269        Ok(())
4270    })()?;
4271    let input_refs = owned_inputs
4272        .iter()
4273        .enumerate()
4274        .map(|(index, tensor)| {
4275            tensor.as_ref().ok_or_else(|| {
4276                Error::Internal(format!(
4277                    "semantic eager JVP derivative input {index} was not populated"
4278                ))
4279            })
4280        })
4281        .collect::<Result<Vec<_>>>()?;
4282    let input_reads = input_refs
4283        .iter()
4284        .map(|tensor| TensorRead::from_tensor(tensor))
4285        .collect::<Vec<_>>();
4286    let outputs = ctx.run_program(engine, &derivative_program, &input_reads)?;
4287    let output_count = outputs.len();
4288    let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
4289        return Err(Error::Internal(format!(
4290            "semantic eager JVP derivative output index {derivative_output_index} is outside {} outputs",
4291            output_count
4292        )));
4293    };
4294    let tangent_trace =
4295        TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&tangent_tensor))?;
4296    let semantic_trace = derivative_trace_from_frozen_program(
4297        &source,
4298        derivative.frozen(),
4299        derivative_output_index,
4300        &[(seed_input_index, Arc::clone(&tangent_tensor))],
4301        &[&output_trace, wrt_trace, &tangent_trace],
4302        None,
4303        "semantic_eager_jvp",
4304    )?;
4305
4306    Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
4307        Arc::clone(ctx),
4308        eager_val_key(),
4309        result,
4310        true,
4311        None,
4312        Some(semantic_trace),
4313    )?)))
4314}
4315
4316fn validate_same_runtime(
4317    runtime: &Arc<EagerRuntime>,
4318    tensor: &EagerTensor,
4319    role: &'static str,
4320) -> Result<()> {
4321    if tensor.ctx_id() != runtime.id() {
4322        return Err(Error::ContextMismatch {
4323            lhs: runtime.id(),
4324            rhs: tensor.ctx_id(),
4325        });
4326    }
4327    let _ = role;
4328    Ok(())
4329}
4330
4331fn copy_value_in_session(session: &mut dyn Engine, value: &RetainedValue) -> Result<Tensor> {
4332    let read = value.tensor_read().map_err(|error| {
4333        Error::runtime_state_source("copy_value_for_runtime", ErrorPhase::Execution, error)
4334    })?;
4335    session.to_contiguous_read(read).map_err(Error::from)
4336}
4337
4338fn validate_seed_tensor(op: &'static str, primal: &EagerTensor, seed: &EagerTensor) -> Result<()> {
4339    if primal.dtype() != seed.dtype() {
4340        return Err(
4341            tenferro_tensor::Error::dtype_mismatch(op, primal.dtype(), seed.dtype()).into(),
4342        );
4343    }
4344    if primal.shape() != seed.shape() {
4345        return Err(
4346            tenferro_tensor::Error::shape_mismatch(op, primal.shape(), seed.shape()).into(),
4347        );
4348    }
4349    Ok(())
4350}
4351
4352/// Eager tensor with reverse-mode autodiff over concrete tensor values.
4353///
4354/// This executes each primitive immediately and records a lightweight reverse
4355/// DAG for `backward()`. Gradients accumulate across repeated `backward()`
4356/// calls until they are cleared explicitly.
4357///
4358/// # Examples
4359///
4360/// ```
4361/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4362///
4363/// let ctx = EagerRuntime::new()?;
4364/// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx)?;
4365/// for _ in 0..2 {
4366///     let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
4367///         let squared = s.mul(&x, &x)?;
4368///         s.reduce_sum(&squared, Some(&[0]))
4369///     })?;
4370///     loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
4371/// }
4372///
4373/// assert_eq!(x.grad()?.unwrap().as_slice::<f64>().unwrap(), &[4.0, 8.0, 12.0]);
4374/// x.clear_grad()?;
4375///
4376/// assert!(x.grad().unwrap().is_none());
4377/// # Ok::<(), tenferro_ad::Error>(())
4378/// ```
4379#[derive(Clone)]
4380pub struct EagerTensor {
4381    pub(crate) key: ValueKey<StdTensorOp>,
4382    pub(crate) trace: Option<EagerTrace>,
4383    pub(crate) semantic_trace: Option<TracedTensor>,
4384    pub(crate) requires_grad: bool,
4385    grad_slot: GradSlot,
4386    pub(crate) ctx: Arc<EagerRuntime>,
4387    _record: Arc<EagerTensorRecord>,
4388}
4389
4390pub(crate) struct EagerTensorRecord {
4391    value: Arc<AdValueRecord>,
4392    key: ValueKey<StdTensorOp>,
4393    trace: Option<EagerTrace>,
4394    semantic_trace: Option<TracedTensor>,
4395    requires_grad: bool,
4396    grad_slot: GradSlot,
4397    ctx: Arc<EagerRuntime>,
4398}
4399
4400struct EagerTensorParts {
4401    ctx: Arc<EagerRuntime>,
4402    key: ValueKey<StdTensorOp>,
4403    requires_grad: bool,
4404    trace: Option<EagerTrace>,
4405    semantic_trace: Option<TracedTensor>,
4406    value: Arc<AdValueRecord>,
4407    register_value: bool,
4408}
4409
4410impl fmt::Debug for EagerTensor {
4411    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
4412        f.debug_struct("EagerTensor")
4413            .field("dtype", &self.dtype())
4414            .field("shape", &self.shape())
4415            .field("key", &self.key)
4416            .field("requires_grad", &self.requires_grad)
4417            .field("has_trace", &self.trace.is_some())
4418            .field("has_semantic_trace", &self.semantic_trace.is_some())
4419            .field("ctx_id", &self.ctx_id())
4420            .finish_non_exhaustive()
4421    }
4422}
4423
4424impl EagerTensor {
4425    /// Create an untracked eager tensor inside an existing eager context.
4426    ///
4427    /// # Examples
4428    ///
4429    /// ```
4430    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4431    ///
4432    /// let ctx = EagerRuntime::new()?;
4433    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
4434    ///
4435    /// assert_eq!(x.value()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
4436    /// # Ok::<(), tenferro_ad::Error>(())
4437    /// ```
4438    ///
4439    /// # Errors
4440    ///
4441    /// Returns [`tenferro_runtime::Error::RuntimeState`] when metadata cannot
4442    /// be registered in the target context, or a typed tensor/backend error
4443    /// while materializing the source value.
4444    pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
4445        Self::new_leaf(ctx, tensor, false)
4446    }
4447
4448    /// Create an untracked eager tensor from compact column-major data inside
4449    /// an existing eager runtime.
4450    ///
4451    /// # Errors
4452    ///
4453    /// Returns [`Error::TensorRuntime`] with
4454    /// [`tenferro_tensor::ValidationError::ShapeMismatch`] when the shape and
4455    /// data length disagree, or with
4456    /// [`tenferro_tensor::ValidationError::IntegerOverflow`] when shape
4457    /// arithmetic overflows. Returns [`Error::RuntimeState`] when eager
4458    /// metadata cannot be registered.
4459    pub fn from_vec_col_major_in<T: TensorScalar>(
4460        shape: impl IntoShapeVec,
4461        data: Vec<T>,
4462        ctx: Arc<EagerRuntime>,
4463    ) -> Result<Self> {
4464        Self::from_tensor_in(Tensor::from_vec_col_major(shape, data)?, ctx)
4465    }
4466
4467    /// Create a tracked eager leaf inside an existing eager context.
4468    ///
4469    /// # Examples
4470    ///
4471    /// ```
4472    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4473    ///
4474    /// let ctx = EagerRuntime::new()?;
4475    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
4476    ///
4477    /// assert!(x.grad().unwrap().is_none());
4478    /// # Ok::<(), tenferro_ad::Error>(())
4479    /// ```
4480    ///
4481    /// # Errors
4482    ///
4483    /// Returns [`tenferro_runtime::Error::RuntimeState`] when gradient metadata
4484    /// cannot be registered in the target context, or a typed tensor/backend
4485    /// error while creating the leaf.
4486    pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
4487        Self::new_leaf(ctx, tensor, true)
4488    }
4489
4490    pub(crate) fn new_leaf(
4491        ctx: Arc<EagerRuntime>,
4492        tensor: Tensor,
4493        requires_grad: bool,
4494    ) -> Result<Self> {
4495        Self::new_leaf_with_session(ctx, tensor, requires_grad, None)
4496    }
4497
4498    pub(crate) fn new_leaf_in_session(
4499        ctx: Arc<EagerRuntime>,
4500        tensor: Tensor,
4501        requires_grad: bool,
4502        session: &mut dyn Engine,
4503    ) -> Result<Self> {
4504        Self::new_leaf_with_session(ctx, tensor, requires_grad, Some(session))
4505    }
4506
4507    fn new_leaf_with_session(
4508        ctx: Arc<EagerRuntime>,
4509        tensor: Tensor,
4510        requires_grad: bool,
4511        session: Option<&mut dyn Engine>,
4512    ) -> Result<Self> {
4513        let key = eager_val_key();
4514        // A host-placement tensor needs no backend session: the CPU backend only
4515        // copies the host buffer for it, so entering a session would add
4516        // admission, provider exclusion, and session construction without doing
4517        // any provider work (#1704). Views and backend-family reads keep the
4518        // session path, and device runtimes keep it for every read.
4519        let read = TensorRead::from_tensor(&tensor);
4520        let semantic_tensor =
4521            match session {
4522                Some(session) => session.to_contiguous_read(read).map_err(Error::from)?,
4523                None => match cpu_host_owned_read(&read) {
4524                    Some(materialized) => materialized.map_err(Error::from)?,
4525                    None => return Err(Error::unsupported(
4526                        "EagerTensor::new_leaf",
4527                        ErrorPhase::Execution,
4528                        "this tensor needs an engine to become an eager leaf; create it through \
4529                         EagerSession::constant_from or EagerSession::variable_from",
4530                    )),
4531                },
4532            };
4533        let semantic_value = Arc::new(RetainedValue::from_tensor(semantic_tensor));
4534        let semantic_trace = TracedTensor::from_shared_tensor_value_symbolic_shape(semantic_value)?;
4535        // Deferred materialization: the per-op/leaf global-metadata registry
4536        // write for the eager tensor key was unreadable after the semantic
4537        // trace became the sole AD carrier, so it is dropped.
4538        // ponytail: leaf input-key metadata is still registered by
4539        // `from_shared_tensor_value_symbolic_shape`; the eager-key entry was
4540        // vestigial and removed. Add back only if something reads it.
4541        let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_leaf")?;
4542        Self::from_parts(EagerTensorParts {
4543            ctx,
4544            key,
4545            requires_grad,
4546            trace: None,
4547            semantic_trace: Some(semantic_trace),
4548            value,
4549            register_value: true,
4550        })
4551    }
4552
4553    pub(crate) fn new_result(
4554        ctx: Arc<EagerRuntime>,
4555        key: ValueKey<StdTensorOp>,
4556        tensor: Tensor,
4557        requires_grad: bool,
4558        trace: Option<EagerTrace>,
4559    ) -> Result<Self> {
4560        Self::new_result_with_semantic_trace(ctx, key, tensor, requires_grad, trace, None)
4561    }
4562
4563    pub(crate) fn new_result_with_semantic_trace(
4564        ctx: Arc<EagerRuntime>,
4565        key: ValueKey<StdTensorOp>,
4566        tensor: Tensor,
4567        requires_grad: bool,
4568        trace: Option<EagerTrace>,
4569        semantic_trace: Option<TracedTensor>,
4570    ) -> Result<Self> {
4571        let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_result")?;
4572        Self::from_parts(EagerTensorParts {
4573            ctx,
4574            key,
4575            requires_grad,
4576            trace,
4577            semantic_trace,
4578            value,
4579            register_value: true,
4580        })
4581    }
4582
4583    pub(crate) fn new_unregistered_result_with_semantic_trace(
4584        ctx: Arc<EagerRuntime>,
4585        key: ValueKey<StdTensorOp>,
4586        tensor: Tensor,
4587        requires_grad: bool,
4588        trace: Option<EagerTrace>,
4589        semantic_trace: Option<TracedTensor>,
4590    ) -> Result<Self> {
4591        let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_unregistered_result")?;
4592        Self::from_parts(EagerTensorParts {
4593            ctx,
4594            key,
4595            requires_grad,
4596            trace,
4597            semantic_trace,
4598            value,
4599            register_value: false,
4600        })
4601    }
4602
4603    pub(crate) fn new_result_value(
4604        ctx: Arc<EagerRuntime>,
4605        key: ValueKey<StdTensorOp>,
4606        value: TensorValue,
4607        requires_grad: bool,
4608        trace: Option<EagerTrace>,
4609        semantic_trace: Option<TracedTensor>,
4610    ) -> Result<Self> {
4611        let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
4612            Error::runtime_state(
4613                "EagerTensor::new_result_value",
4614                ErrorPhase::Execution,
4615                "a TensorValue could not be transferred into its allocation group",
4616            )
4617        })?;
4618        let value = AdValueRecord::from_group(group, slot, dtype, shape);
4619        Self::from_parts(EagerTensorParts {
4620            ctx,
4621            key,
4622            requires_grad,
4623            trace,
4624            semantic_trace,
4625            value,
4626            register_value: true,
4627        })
4628    }
4629
4630    fn from_parts(parts: EagerTensorParts) -> Result<Self> {
4631        let EagerTensorParts {
4632            ctx,
4633            key,
4634            requires_grad,
4635            trace,
4636            semantic_trace,
4637            value,
4638            register_value,
4639        } = parts;
4640        let grad_slot = Arc::new(Mutex::new(None));
4641        if requires_grad {
4642            ctx.try_register_grad_slot(&key, &grad_slot)?;
4643        }
4644        let record = Arc::new(EagerTensorRecord {
4645            value: Arc::clone(&value),
4646            key: key.clone(),
4647            trace: trace.clone(),
4648            semantic_trace: semantic_trace.clone(),
4649            requires_grad,
4650            grad_slot: Arc::clone(&grad_slot),
4651            ctx: Arc::clone(&ctx),
4652        });
4653        if register_value {
4654            ctx.try_register_value_record(&key, &record)?;
4655        }
4656
4657        Ok(Self {
4658            key,
4659            trace,
4660            semantic_trace,
4661            requires_grad,
4662            grad_slot,
4663            ctx,
4664            _record: record,
4665        })
4666    }
4667
4668    pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
4669        let value =
4670            AdValueRecord::from_untracked_tensor(tensor, "EagerTensor::new_untracked_result")?;
4671        Ok(Self::new_untracked_value_record(ctx, value, None))
4672    }
4673
4674    pub(crate) fn new_untracked_value_result(
4675        ctx: Arc<EagerRuntime>,
4676        value: TensorValue,
4677    ) -> Result<Self> {
4678        Self::new_untracked_value_result_with_semantic_trace(ctx, value, None)
4679    }
4680
4681    pub(crate) fn new_untracked_value_result_with_semantic_trace(
4682        ctx: Arc<EagerRuntime>,
4683        value: TensorValue,
4684        semantic_trace: Option<TracedTensor>,
4685    ) -> Result<Self> {
4686        let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
4687            Error::runtime_state(
4688                "EagerTensor::new_untracked_value_result",
4689                ErrorPhase::Execution,
4690                "a TensorValue could not be transferred into its allocation group",
4691            )
4692        })?;
4693        let value = AdValueRecord::from_group(group, slot, dtype, shape);
4694        Ok(Self::new_untracked_value_record(ctx, value, semantic_trace))
4695    }
4696
4697    fn new_untracked_value_record(
4698        ctx: Arc<EagerRuntime>,
4699        value: Arc<AdValueRecord>,
4700        semantic_trace: Option<TracedTensor>,
4701    ) -> Self {
4702        let key = eager_val_key();
4703        let grad_slot = Arc::new(Mutex::new(None));
4704        let record = Arc::new(EagerTensorRecord {
4705            value,
4706            key: key.clone(),
4707            trace: None,
4708            semantic_trace: semantic_trace.clone(),
4709            requires_grad: false,
4710            grad_slot: Arc::clone(&grad_slot),
4711            ctx: Arc::clone(&ctx),
4712        });
4713        Self {
4714            key,
4715            trace: None,
4716            semantic_trace,
4717            requires_grad: false,
4718            grad_slot,
4719            ctx,
4720            _record: record,
4721        }
4722    }
4723
4724    pub(crate) fn from_record(record: Arc<EagerTensorRecord>) -> Self {
4725        Self {
4726            key: record.key.clone(),
4727            trace: record.trace.clone(),
4728            semantic_trace: record.semantic_trace.clone(),
4729            requires_grad: record.requires_grad,
4730            grad_slot: Arc::clone(&record.grad_slot),
4731            ctx: Arc::clone(&record.ctx),
4732            _record: record,
4733        }
4734    }
4735
4736    /// Detach this tensor from the reverse graph.
4737    ///
4738    /// The returned tensor keeps the concrete value but no longer contributes
4739    /// gradients to the original graph.
4740    ///
4741    /// # Errors
4742    ///
4743    /// Returns [`Error::TensorRuntime`] when `engine` cannot materialize the
4744    /// retained value, [`Error::RuntimeStateSource`] when the value has no
4745    /// readable tensor or symbolic graph metadata registration is unavailable,
4746    /// and [`Error::RuntimeState`] when the detached record's allocation-group
4747    /// descriptor is invalid.
4748    ///
4749    /// # Examples
4750    ///
4751    /// ```
4752    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4753    ///
4754    /// let ctx = EagerRuntime::new()?;
4755    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
4756    /// let y = x.detach(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
4757    ///
4758    /// assert_eq!(y.value()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
4759    /// assert!(y.grad().unwrap().is_none());
4760    /// # Ok::<(), tenferro_ad::Error>(())
4761    /// ```
4762    pub fn detach(&self, engine: &mut dyn Engine) -> Result<Self> {
4763        let tensor = self.duplicate_value(engine)?;
4764        let semantic_trace = TracedTensor::from_tensor_symbolic_shape(tensor)?;
4765        Ok(Self::new_untracked_value_record(
4766            self.ctx.clone(),
4767            Arc::clone(&self._record.value),
4768            Some(semantic_trace),
4769        ))
4770    }
4771
4772    /// Detach this tensor from its graph and re-register it in a different
4773    /// context as an untracked leaf.
4774    ///
4775    /// # Examples
4776    ///
4777    /// ```
4778    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4779    ///
4780    /// let ctx_a = EagerRuntime::new()?;
4781    /// let ctx_b = EagerRuntime::new()?;
4782    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx_a)?;
4783    /// let d = x.detach_into(&mut tenferro_cpu::CpuEngine::serial(), &ctx_b)?;
4784    ///
4785    /// assert!(!d.tracks_grad());
4786    /// assert_eq!(d.ctx_id(), ctx_b.id());
4787    /// # Ok::<(), tenferro_ad::Error>(())
4788    /// ```
4789    ///
4790    /// # Errors
4791    ///
4792    /// Returns [`Error::RuntimeState`] if the source cannot be materialized or
4793    /// the target context cannot register its metadata.
4794    pub fn detach_into(&self, engine: &mut dyn Engine, ctx: &Arc<EagerRuntime>) -> Result<Self> {
4795        let tensor = self.to_tensor(&mut *engine)?;
4796        Self::new_leaf_in_session(Arc::clone(ctx), tensor, false, engine)
4797    }
4798
4799    /// Borrow the retained value without creating an owner or copy.
4800    ///
4801    /// # Errors
4802    ///
4803    /// Returns [`Error::RuntimeState`] when the retained allocation-group
4804    /// descriptor is unavailable or invalid.
4805    pub fn value(&self) -> Result<ValueGuard<'_>> {
4806        self._record.value.value("EagerTensor::value")
4807    }
4808
4809    /// Explicitly duplicate this value into a fresh standalone allocation.
4810    ///
4811    /// # Errors
4812    ///
4813    /// Returns [`Error::RuntimeState`] when the retained value or execution
4814    /// session is unavailable, or a typed host/backend error when the value
4815    /// cannot be materialized as a contiguous tensor.
4816    ///
4817    /// # Examples
4818    ///
4819    /// ```
4820    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4821    ///
4822    /// let ctx = EagerRuntime::new()?;
4823    /// let value = EagerTensor::from_tensor_in(
4824    ///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?,
4825    ///     ctx,
4826    /// )?;
4827    /// let duplicate = value.duplicate_value(&mut tenferro_cpu::CpuEngine::serial())?;
4828    /// assert_eq!(duplicate.as_slice::<f64>()?, &[1.0, 2.0]);
4829    /// # Ok::<(), tenferro_ad::Error>(())
4830    /// ```
4831    pub fn duplicate_value(&self, engine: &mut dyn Engine) -> Result<Tensor> {
4832        self.duplicate_value_in_session(engine)
4833    }
4834
4835    fn duplicate_host_value(&self) -> Option<Tensor> {
4836        // A pooled value duplicates through its descriptor view; a caller-owned
4837        // payload has no such view and duplicates through its own read path, which
4838        // copies the value while keeping its element type.
4839        self.value().ok()?.duplicate_host_tensor().ok()
4840    }
4841
4842    pub(crate) fn duplicate_value_in_session(&self, session: &mut dyn Engine) -> Result<Tensor> {
4843        if let Some(tensor) = self.duplicate_host_value() {
4844            return Ok(tensor);
4845        }
4846        let read = self
4847            ._record
4848            .value
4849            .tensor_read("EagerTensor::duplicate_value")?;
4850        session.to_contiguous_read(read).map_err(Error::from)
4851    }
4852
4853    // INVARIANT: the error variants return the unchanged eager handle so a
4854    // caller can retry ownership extraction without an implicit copy.
4855    #[allow(clippy::result_large_err)]
4856    /// Consume this handle and structurally extract its retained allocation.
4857    ///
4858    /// A shared handle is returned unchanged as [`IntoValueError::NotUnique`].
4859    /// Group extraction failures return the unchanged handle and typed group
4860    /// error; no copy or fallback materialization is attempted.
4861    ///
4862    /// # Errors
4863    ///
4864    /// Returns [`IntoValueError::NotUnique`] when another handle retains the
4865    /// value, or [`IntoValueError::Extract`] when structural group extraction
4866    /// fails because the allocation is aliased or its descriptor is invalid.
4867    ///
4868    /// # Examples
4869    ///
4870    /// ```
4871    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4872    ///
4873    /// let ctx = EagerRuntime::new()?;
4874    /// let value = EagerTensor::from_tensor_in(
4875    ///     Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?,
4876    ///     ctx,
4877    /// )?;
4878    /// let owner = value
4879    ///     .into_value()
4880    ///     .expect("a uniquely owned value should be extractable");
4881    /// assert_eq!(owner.as_slice::<f64>()?, &[3.0]);
4882    /// # Ok::<(), tenferro_ad::Error>(())
4883    /// ```
4884    pub fn into_value(self) -> std::result::Result<Tensor, IntoValueError<Self>> {
4885        if Arc::strong_count(&self._record) != 1 {
4886            return Err(IntoValueError::NotUnique(self));
4887        }
4888        let Self { _record, .. } = self;
4889        let record = match Arc::try_unwrap(_record) {
4890            Ok(record) => record,
4891            Err(record) => return Err(IntoValueError::NotUnique(Self::from_record(record))),
4892        };
4893        let EagerTensorRecord {
4894            value,
4895            key,
4896            trace,
4897            semantic_trace,
4898            requires_grad,
4899            grad_slot,
4900            ctx,
4901        } = record;
4902        let value = match Arc::try_unwrap(value) {
4903            Ok(value) => value,
4904            Err(value) => {
4905                let record = Arc::new(EagerTensorRecord {
4906                    value,
4907                    key,
4908                    trace,
4909                    semantic_trace,
4910                    requires_grad,
4911                    grad_slot,
4912                    ctx,
4913                });
4914                return Err(IntoValueError::NotUnique(Self::from_record(record)));
4915            }
4916        };
4917        let AdValueRecord {
4918            container,
4919            dtype,
4920            shape,
4921        } = value;
4922        let container = match Arc::try_unwrap(container) {
4923            Ok(container) => container,
4924            Err(container) => {
4925                let record = Arc::new(EagerTensorRecord {
4926                    value: Arc::new(AdValueRecord {
4927                        container,
4928                        dtype,
4929                        shape,
4930                    }),
4931                    key,
4932                    trace,
4933                    semantic_trace,
4934                    requires_grad,
4935                    grad_slot,
4936                    ctx,
4937                });
4938                return Err(IntoValueError::NotUnique(Self::from_record(record)));
4939            }
4940        };
4941        let (group, slot) = match container {
4942            // A caller-owned payload is handed back to its owner unchanged.
4943            RetentionContainer::CallerOwned { tensor } => return Ok(*tensor),
4944            RetentionContainer::Owned { tensor } => return Ok(tensor),
4945            RetentionContainer::Pooled { group, slot } => (group, slot),
4946        };
4947        match group.into_tensor(slot) {
4948            Ok(tensor) => Ok(tensor),
4949            Err((group, error)) => {
4950                let record = Arc::new(EagerTensorRecord {
4951                    value: Arc::new(AdValueRecord {
4952                        container: Arc::new(RetentionContainer::Pooled { group, slot }),
4953                        dtype,
4954                        shape,
4955                    }),
4956                    key,
4957                    trace,
4958                    semantic_trace,
4959                    requires_grad,
4960                    grad_slot,
4961                    ctx,
4962                });
4963                Err(IntoValueError::Extract {
4964                    value: Self::from_record(record),
4965                    error,
4966                })
4967            }
4968        }
4969    }
4970
4971    /// Return this tensor's scalar dtype without materializing through
4972    /// [`value`](Self::value).
4973    pub fn dtype(&self) -> DType {
4974        self._record.value.dtype()
4975    }
4976
4977    /// Return this tensor's logical shape without materializing through
4978    /// [`value`](Self::value).
4979    pub fn shape(&self) -> &[usize] {
4980        self._record.value.shape()
4981    }
4982
4983    /// Borrow this tensor value as a [`TensorRead`].
4984    ///
4985    /// This is the preferred borrowed input boundary for executor calls. It
4986    /// preserves the option to replace eager storage with non-contiguous views
4987    /// without forcing callers through [`value`](Self::value).
4988    ///
4989    /// # Panics
4990    ///
4991    /// Panics if a validated eager value record becomes unavailable, which
4992    /// indicates an internal invariant violation.
4993    pub fn tensor_read(&self) -> TensorRead<'_> {
4994        self._record
4995            .value
4996            .tensor_read("EagerTensor::tensor_read")
4997            .expect("validated eager value record")
4998    }
4999
5000    /// Materialize this eager tensor as an owned [`Tensor`].
5001    ///
5002    /// This is the owned materialization boundary for callers that need a
5003    /// standalone compact tensor. The operation is fallible because eager
5004    /// values may be backed by lazy or backend-resident storage.
5005    ///
5006    /// # Errors
5007    ///
5008    /// Returns [`Error::RuntimeState`] if backend state is unavailable, or a
5009    /// typed tensor backend error when contiguous materialization fails.
5010    pub fn to_tensor(&self, engine: &mut dyn Engine) -> Result<Tensor> {
5011        self.duplicate_value(engine)
5012    }
5013
5014    /// Return the accumulated gradient currently stored for this tensor.
5015    ///
5016    /// The stored gradient accumulates across repeated `backward()` calls
5017    /// until it is cleared explicitly.
5018    ///
5019    /// For complex scalar losses, stored gradients use tenferro's
5020    /// Hermitian-adjoint cotangent convention. See
5021    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
5022    ///
5023    /// # Examples
5024    ///
5025    /// ```
5026    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5027    ///
5028    /// let ctx = EagerRuntime::new()?;
5029    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
5030    /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
5031    ///     let y = s.exp(&x)?;
5032    ///     s.reduce_sum(&y, Some(&[0]))
5033    /// })?;
5034    /// let _cotangents = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
5035    ///
5036    /// let grad = x.grad()?.unwrap();
5037    /// assert_eq!(grad.shape(), &[2]);
5038    /// # Ok::<(), tenferro_ad::Error>(())
5039    /// ```
5040    ///
5041    /// # Errors
5042    ///
5043    /// Returns [`Error::RuntimeState`] if the gradient slot is poisoned or no
5044    /// longer available.
5045    pub fn grad(&self) -> Result<Option<GradientValue>> {
5046        self.grad_slot
5047            .lock()
5048            .map_err(|_| {
5049                Error::runtime_state(
5050                    "eager_gradient_slot",
5051                    ErrorPhase::Execution,
5052                    "lock poisoned",
5053                )
5054            })
5055            .map(|slot| {
5056                slot.as_ref().map(|record| GradientValue {
5057                    record: Arc::clone(record),
5058                })
5059            })
5060    }
5061
5062    /// Clear the accumulated gradient stored for this tensor.
5063    ///
5064    /// This only affects this tensor's gradient slot. Other tensors in the
5065    /// same context retain their gradients until they are cleared explicitly or
5066    /// overwritten by later accumulation.
5067    ///
5068    /// # Examples
5069    ///
5070    /// ```
5071    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5072    ///
5073    /// let ctx = EagerRuntime::new()?;
5074    /// 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();
5075    /// let y = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![4.0_f64, 5.0, 6.0]).unwrap(), ctx).unwrap();
5076    /// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
5077    ///     let product = s.mul(&x, &y)?;
5078    ///     s.reduce_sum(&product, Some(&[0]))
5079    /// })?;
5080    /// let _ = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
5081    ///
5082    /// x.clear_grad()?;
5083    ///
5084    /// assert!(x.grad()?.is_none());
5085    /// assert!(y.grad()?.is_some());
5086    /// # Ok::<(), tenferro_ad::Error>(())
5087    /// ```
5088    ///
5089    /// # Errors
5090    ///
5091    /// Returns [`Error::RuntimeState`] if the gradient slot lock is poisoned.
5092    pub fn clear_grad(&self) -> Result<()> {
5093        *self.grad_slot.lock().map_err(|_| {
5094            Error::runtime_state(
5095                "eager_gradient_slot",
5096                ErrorPhase::Execution,
5097                "lock poisoned",
5098            )
5099        })? = None;
5100        Ok(())
5101    }
5102
5103    /// Report whether this tensor participates in gradient tracking.
5104    ///
5105    /// Tracked tensors keep a gradient slot in their eager context; untracked
5106    /// tensors and detached tensors do not.
5107    ///
5108    /// # Examples
5109    ///
5110    /// ```
5111    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5112    ///
5113    /// let ctx = EagerRuntime::new()?;
5114    /// let plain = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
5115    /// let tracked = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
5116    /// let detached = tracked.detach(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
5117    ///
5118    /// assert!(!plain.tracks_grad());
5119    /// assert!(tracked.tracks_grad());
5120    /// assert!(!detached.tracks_grad());
5121    /// # Ok::<(), tenferro_ad::Error>(())
5122    /// ```
5123    pub fn tracks_grad(&self) -> bool {
5124        self.requires_grad
5125    }
5126
5127    #[cfg(test)]
5128    fn debug_trace_saved_value_count(&self) -> Option<usize> {
5129        None
5130    }
5131
5132    /// Return the opaque identifier of the context this tensor belongs to.
5133    ///
5134    /// # Examples
5135    ///
5136    /// ```
5137    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5138    ///
5139    /// let ctx = EagerRuntime::new()?;
5140    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
5141    ///
5142    /// assert_eq!(x.ctx_id(), ctx.id());
5143    /// # Ok::<(), tenferro_ad::Error>(())
5144    /// ```
5145    pub fn ctx_id(&self) -> ContextId {
5146        self.ctx.id()
5147    }
5148
5149    /// Borrow the eager runtime context that owns this tensor.
5150    pub fn runtime(&self) -> &Arc<EagerRuntime> {
5151        &self.ctx
5152    }
5153
5154    /// Check whether two tensors belong to the same eager context.
5155    ///
5156    /// # Examples
5157    ///
5158    /// ```
5159    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5160    ///
5161    /// let ctx = EagerRuntime::new()?;
5162    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
5163    /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx).unwrap();
5164    ///
5165    /// assert!(x.same_context(&y));
5166    /// # Ok::<(), tenferro_ad::Error>(())
5167    /// ```
5168    pub fn same_context(&self, other: &Self) -> bool {
5169        self.ctx_id() == other.ctx_id()
5170    }
5171
5172    #[cfg(test)]
5173    pub(crate) fn standard_graph_op(
5174        engine: &mut dyn ExecutionEngine,
5175        inputs: &[&Self],
5176        build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
5177    ) -> Result<Vec<Self>> {
5178        let Some(first) = inputs.first() else {
5179            return Err(Error::Internal(
5180                "standard eager graph op requires at least one input tensor".to_string(),
5181            ));
5182        };
5183        let ctx = Arc::clone(&first.ctx);
5184        for tensor in inputs.iter().skip(1) {
5185            if !first.same_context(tensor) {
5186                return Err(Error::ContextMismatch {
5187                    lhs: first.ctx_id(),
5188                    rhs: tensor.ctx_id(),
5189                });
5190            }
5191        }
5192
5193        let graph_input_keys = (0..inputs.len())
5194            .map(|_| next_input_key())
5195            .collect::<Vec<_>>();
5196        let graph = build_graph(&graph_input_keys)?;
5197        let initial_data = graph_input_keys
5198            .iter()
5199            .zip(inputs.iter())
5200            .map(|(key, tensor)| {
5201                Ok((
5202                    ValueKey::Input(key.clone()),
5203                    tensor.to_tensor(&mut *engine)?,
5204                ))
5205            })
5206            .collect::<Result<HashMap<_, _>>>()?;
5207        let execution = ctx.exec_standard_graph_outputs(engine, graph.as_ref(), initial_data)?;
5208        if execution.outputs.len() != graph.outputs().len() {
5209            return Err(Error::Internal(format!(
5210                "standard eager graph op expected {} graph outputs, got {}",
5211                graph.outputs().len(),
5212                execution.outputs.len()
5213            )));
5214        }
5215
5216        if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
5217            return execution
5218                .outputs
5219                .into_iter()
5220                .map(|output| {
5221                    Self::new_unregistered_result_with_semantic_trace(
5222                        Arc::clone(&ctx),
5223                        eager_val_key(),
5224                        output,
5225                        false,
5226                        None,
5227                        None,
5228                    )
5229                })
5230                .collect();
5231        }
5232
5233        let recorded = record_eager_graph_outputs(
5234            graph.as_ref(),
5235            &graph_input_keys,
5236            &execution.outputs,
5237            inputs,
5238        )?;
5239        if recorded.traces.len() != execution.outputs.len() {
5240            return Err(Error::Internal(format!(
5241                "standard eager graph op expected {} eager traces, got {}",
5242                execution.outputs.len(),
5243                recorded.traces.len()
5244            )));
5245        }
5246
5247        recorded
5248            .traces
5249            .into_iter()
5250            .zip(recorded.semantic_traces)
5251            .zip(execution.outputs)
5252            .map(|((trace, semantic_trace), output)| {
5253                Self::new_result_with_semantic_trace(
5254                    Arc::clone(&ctx),
5255                    trace.key,
5256                    output,
5257                    trace.requires_grad,
5258                    trace.trace,
5259                    semantic_trace,
5260                )
5261            })
5262            .collect()
5263    }
5264
5265    /// Run reverse-mode AD from this scalar output.
5266    ///
5267    /// Returns the full cotangent map produced by the reverse pass and also
5268    /// accumulates into `grad()` for tracked eager tensors reachable from this
5269    /// output.
5270    ///
5271    /// For complex scalar outputs, cotangents use tenferro's Hermitian
5272    /// real-inner-product convention. See
5273    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
5274    ///
5275    /// # Examples
5276    ///
5277    /// ```
5278    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5279    ///
5280    /// let ctx = EagerRuntime::new()?;
5281    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx).unwrap();
5282    /// for _ in 0..2 {
5283    ///     let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
5284    ///         let doubled = s.add(&x, &x)?;
5285    ///         s.reduce_sum(&doubled, Some(&[0]))
5286    ///     })?;
5287    ///     loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
5288    /// }
5289    ///
5290    /// assert_eq!(x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(), &[4.0, 4.0, 4.0]);
5291    /// # Ok::<(), tenferro_ad::Error>(())
5292    /// ```
5293    ///
5294    /// # Errors
5295    ///
5296    /// Returns [`Error::NonScalarGrad`] when this output is not scalar,
5297    /// [`Error::UnsupportedAdRule`] when a graph operation lacks a reverse rule,
5298    /// or a typed validation/backend/runtime-state error during the reverse pass.
5299    pub fn backward(&self, engine: &mut dyn ExecutionEngine) -> Result<Gradients> {
5300        if !self.shape().is_empty() {
5301            return Err(Error::NonScalarGrad {
5302                shape: self.shape().to_vec(),
5303            });
5304        }
5305
5306        let value = self.to_tensor(&mut *engine)?;
5307        let seed = one_like_tensor(&value, &mut *engine)?;
5308        self.backward_from_seed(engine, seed)
5309    }
5310
5311    /// Run reverse-mode AD from this output with an explicit cotangent seed.
5312    ///
5313    /// This is the stateful eager VJP sugar: it returns the cotangent map and
5314    /// accumulates reachable tracked leaves into their `grad()` slots. Use
5315    /// [`EagerRuntime::vjp`] when the VJP result should be returned as a
5316    /// composable eager tensor without touching grad slots.
5317    ///
5318    /// # Examples
5319    ///
5320    /// ```
5321    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5322    ///
5323    /// let ctx = EagerRuntime::new()?;
5324    /// let x = EagerTensor::requires_grad_in(
5325    ///     Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
5326    ///     ctx.clone(),
5327    /// )?;
5328    /// let seed = EagerTensor::from_tensor_in(
5329    ///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
5330    ///     ctx,
5331    /// )?;
5332    /// let y = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
5333    /// y.backward_with(&mut tenferro_cpu::CpuEngine::serial(), &seed)?;
5334    /// assert_eq!(x.grad()?.unwrap().as_slice::<f64>().unwrap(), &[4.0, 12.0]);
5335    /// # Ok::<(), tenferro_ad::Error>(())
5336    /// ```
5337    ///
5338    /// # Errors
5339    ///
5340    /// Returns [`Error::ContextMismatch`] when `cotangent` belongs to another
5341    /// eager runtime, [`Error::Validation`] when its shape or dtype is not a
5342    /// valid seed, [`Error::UnsupportedAdRule`] for an unavailable reverse
5343    /// rule, or a typed backend/runtime-state error during execution.
5344    pub fn backward_with(
5345        &self,
5346        engine: &mut dyn ExecutionEngine,
5347        cotangent: &EagerTensor,
5348    ) -> Result<Gradients> {
5349        if !self.same_context(cotangent) {
5350            return Err(Error::ContextMismatch {
5351                lhs: self.ctx_id(),
5352                rhs: cotangent.ctx_id(),
5353            });
5354        }
5355        validate_seed_tensor("backward", self, cotangent)?;
5356        let seed = cotangent.to_tensor(&mut *engine)?;
5357        self.backward_from_seed(engine, seed)
5358    }
5359
5360    fn backward_from_seed(
5361        &self,
5362        engine: &mut dyn ExecutionEngine,
5363        seed: Tensor,
5364    ) -> Result<Gradients> {
5365        let cotangent =
5366            EagerTensor::new_result(Arc::clone(&self.ctx), eager_val_key(), seed, false, None)?;
5367        let candidate_keys = {
5368            let mut slots = self.ctx.lock_grad_slots()?;
5369            let mut keys = Vec::new();
5370            slots.retain(|key, slot| {
5371                if slot.upgrade().is_some() {
5372                    keys.push(key.clone());
5373                    true
5374                } else {
5375                    false
5376                }
5377            });
5378            keys
5379        };
5380
5381        let mut targets = Vec::new();
5382        for key in candidate_keys {
5383            let Some(record) = self.ctx.value_record(&key)? else {
5384                continue;
5385            };
5386            if !record.requires_grad {
5387                continue;
5388            }
5389            targets.push((key, EagerTensor::from_record(record)));
5390        }
5391        let wrts = targets.iter().map(|(_, wrt)| wrt).collect::<Vec<_>>();
5392        let gradients = semantic_eager_vjp_many(&self.ctx, engine, self, &wrts, &cotangent)?;
5393        let cotangents = (|| {
5394            let session: &mut dyn Engine = &mut *engine;
5395            let mut cotangents = HashMap::new();
5396            for ((key, _), grad) in targets.into_iter().zip(gradients) {
5397                let Some(grad) = grad else {
5398                    continue;
5399                };
5400                let tensor = match grad.into_value() {
5401                    Ok(tensor) => tensor,
5402                    Err(IntoValueError::NotUnique(handle)) => {
5403                        handle.duplicate_value_in_session(session)?
5404                    }
5405                    // A gradient whose retained layout is a view (for example a
5406                    // transpose) cannot be extracted as an owned tensor; copy
5407                    // it to a compact tensor instead.
5408                    Err(IntoValueError::Extract { value, .. }) => {
5409                        value.duplicate_value_in_session(session)?
5410                    }
5411                };
5412                cotangents.insert(key, tensor);
5413            }
5414            self.ctx.store_grads(&cotangents, session)?;
5415            Ok::<_, Error>(cotangents)
5416        })()?;
5417        Gradients::from_tensors(cotangents)
5418    }
5419}
5420
5421/// Insert a weak registry entry, first dropping dead entries when the insert
5422/// would otherwise grow the table.
5423///
5424/// Dead entries are otherwise removed only when their own key is looked up, so
5425/// a long-running runtime would keep one entry per value it ever created.
5426/// Pruning at the growth point keeps the cost amortized O(1) per insert: a
5427/// sweep runs at most once per capacity doubling.
5428fn insert_pruning_dead<K: std::hash::Hash + Eq, V>(
5429    map: &mut HashMap<K, Weak<V>>,
5430    key: K,
5431    value: Weak<V>,
5432) {
5433    if map.len() == map.capacity() {
5434        map.retain(|_, entry| entry.strong_count() > 0);
5435    }
5436    map.insert(key, value);
5437}
5438
5439pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
5440    ValueKey::Input(next_input_key())
5441}
5442
5443pub(crate) struct RecordedEagerTrace {
5444    pub(crate) key: ValueKey<StdTensorOp>,
5445    pub(crate) trace: Option<EagerTrace>,
5446    pub(crate) requires_grad: bool,
5447}
5448
5449pub(crate) struct RecordedEagerOutputs {
5450    pub(crate) traces: Vec<RecordedEagerTrace>,
5451    pub(crate) semantic_traces: Vec<Option<TracedTensor>>,
5452}
5453
5454pub(crate) fn record_eager_outputs_in_session(
5455    op: &StdTensorOp,
5456    outputs: &[&Tensor],
5457    inputs: &[&EagerTensor],
5458    session: &mut dyn Engine,
5459) -> Result<RecordedEagerOutputs> {
5460    let metadata = outputs
5461        .iter()
5462        .map(|output| tensor_meta_from_tensor(output))
5463        .collect();
5464    record_eager_outputs_inner(op, metadata, inputs, Some(session))
5465}
5466
5467pub(crate) fn record_eager_value_outputs_in_session(
5468    op: &StdTensorOp,
5469    outputs: &[&TensorValue],
5470    inputs: &[&EagerTensor],
5471    session: &mut dyn Engine,
5472) -> Result<RecordedEagerOutputs> {
5473    let metadata = outputs
5474        .iter()
5475        .map(|output| tensor_meta_from_value(output))
5476        .collect();
5477    record_eager_outputs_inner(op, metadata, inputs, Some(session))
5478}
5479
5480fn record_eager_outputs_inner(
5481    op: &StdTensorOp,
5482    output_metadata: Vec<TensorMeta>,
5483    inputs: &[&EagerTensor],
5484    session: Option<&mut dyn Engine>,
5485) -> Result<RecordedEagerOutputs> {
5486    let semantic_traces = record_semantic_eager_outputs(op, &output_metadata, inputs, session)?;
5487    record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
5488}
5489
5490fn record_semantic_eager_outputs(
5491    op: &StdTensorOp,
5492    output_metadata: &[TensorMeta],
5493    inputs: &[&EagerTensor],
5494    mut session: Option<&mut dyn Engine>,
5495) -> Result<Vec<Option<TracedTensor>>> {
5496    // Materialize a constant semantic leaf for any untracked input that lost
5497    // its implicit semantic trace on the active-edge fast path. This keeps
5498    // "untracked constant feeds tracked AD" working (PyTorch-style: untracked
5499    // = constant leaf, no gradient flows to it) without re-recording every
5500    // untracked op at creation time.
5501    let mut owned_constants = Vec::<TracedTensor>::new();
5502    for input in inputs {
5503        if input.semantic_trace.is_none() {
5504            let value = match session.as_deref_mut() {
5505                Some(session) => input.duplicate_value_in_session(session)?,
5506                None => input.duplicate_host_value().ok_or_else(|| {
5507                    Error::unsupported(
5508                        "record_semantic_eager_outputs",
5509                        ErrorPhase::Execution,
5510                        "a non-host untracked input needs an engine to become a constant leaf",
5511                    )
5512                })?,
5513            };
5514            owned_constants.push(TracedTensor::from_tensor_symbolic_shape(value)?);
5515        }
5516    }
5517    let mut constants = owned_constants.iter();
5518    let mut semantic_inputs: Vec<&TracedTensor> = inputs
5519        .iter()
5520        .map(|input| {
5521            input
5522                .semantic_trace
5523                .as_ref()
5524                .unwrap_or_else(|| constants.next().expect("materialized constant"))
5525        })
5526        .collect();
5527    let promotion_plan =
5528        eager_input_promotion_plan(op, inputs.len(), |index| inputs[index].dtype());
5529    // Mirror eager execution in the deferred carrier only. The concrete
5530    // tensors have already been promoted at the execution boundary, so these
5531    // casts add semantic graph nodes without an eager copy or backend kernel.
5532    let promoted_semantic_inputs = if semantic_inputs.iter().enumerate().any(|(index, semantic)| {
5533        semantic.dtype != promotion_plan.target_dtype(index, semantic.dtype)
5534    }) {
5535        Some(
5536            semantic_inputs
5537                .iter()
5538                .enumerate()
5539                .map(|(index, &semantic)| {
5540                    let target = promotion_plan.target_dtype(index, semantic.dtype);
5541                    if semantic.dtype == target {
5542                        Ok(Cow::Borrowed(semantic))
5543                    } else {
5544                        semantic.cast(target).map(Cow::Owned)
5545                    }
5546                })
5547                .collect::<Result<Vec<Cow<'_, TracedTensor>>>>()?,
5548        )
5549    } else {
5550        None
5551    };
5552    if let Some(promoted_semantic_inputs) = &promoted_semantic_inputs {
5553        semantic_inputs = promoted_semantic_inputs.iter().map(Cow::as_ref).collect();
5554    }
5555    let exact_semantic_inputs = if matches!(op, StdTensorOp::Concatenate { .. }) {
5556        Some(
5557            semantic_inputs
5558                .iter()
5559                .zip(inputs)
5560                .map(|(&semantic, input)| {
5561                    if semantic.is_concrete_shape() {
5562                        Ok(semantic.clone())
5563                    } else {
5564                        semantic.reshape(input.shape())
5565                    }
5566                })
5567                .collect::<Result<Vec<_>>>()?,
5568        )
5569    } else {
5570        None
5571    };
5572    if let Some(exact_semantic_inputs) = &exact_semantic_inputs {
5573        semantic_inputs = exact_semantic_inputs.iter().collect();
5574    }
5575    // Deferred materialization (issue #1665 steps 6-7): append only a raw
5576    // carrier. The runtime helper retains metadata scopes introduced by the
5577    // promotion/exactification helpers without analyzing this operation.
5578    let outputs = tenferro_runtime::extension::append_raw_eager_outputs(
5579        op.clone(),
5580        &semantic_inputs,
5581        output_metadata,
5582    )?;
5583    Ok(outputs.into_iter().map(Some).collect())
5584}
5585
5586#[cfg(test)]
5587fn record_eager_graph_outputs(
5588    graph: &Graph<StdTensorOp>,
5589    graph_input_keys: &[TensorInputKey],
5590    outputs: &[Tensor],
5591    inputs: &[&EagerTensor],
5592) -> Result<RecordedEagerOutputs> {
5593    let semantic_traces = record_semantic_eager_graph_outputs(graph, graph_input_keys, inputs)?;
5594    let output_metadata = outputs.iter().map(tensor_meta_from_tensor);
5595    record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
5596}
5597
5598#[cfg(test)]
5599fn record_semantic_eager_graph_outputs(
5600    graph: &Graph<StdTensorOp>,
5601    graph_input_keys: &[TensorInputKey],
5602    inputs: &[&EagerTensor],
5603) -> Result<Vec<Option<TracedTensor>>> {
5604    let Some(semantic_inputs) = inputs
5605        .iter()
5606        .map(|input| input.semantic_trace.as_ref())
5607        .collect::<Option<Vec<_>>>()
5608    else {
5609        return Ok(vec![None; graph.outputs().len()]);
5610    };
5611    if graph_input_keys.len() != semantic_inputs.len() {
5612        return Err(Error::Internal(format!(
5613            "semantic graph recording expected {} input keys, got {}",
5614            semantic_inputs.len(),
5615            graph_input_keys.len()
5616        )));
5617    }
5618
5619    let mut values = HashMap::new();
5620    for (key, tensor) in graph_input_keys.iter().zip(semantic_inputs) {
5621        values.insert(ValueKey::Input(key.clone()), tensor.clone());
5622    }
5623
5624    for op_node in graph.operations() {
5625        let input_values = op_node
5626            .inputs
5627            .iter()
5628            .map(|input| {
5629                let key = match input {
5630                    ValueRef::Local(local_id) => &graph.values()[*local_id].key,
5631                    ValueRef::External(key) => key,
5632                };
5633                values.get(key).cloned().ok_or_else(|| {
5634                    Error::Internal(format!(
5635                        "semantic graph recording missing value for {key:?}"
5636                    ))
5637                })
5638            })
5639            .collect::<Result<Vec<_>>>()?;
5640        let input_refs = input_values.iter().collect::<Vec<_>>();
5641        let semantic_outputs = match &op_node.operation {
5642            StdTensorOp::Extension(ext) => {
5643                tenferro_runtime::extension::apply(Arc::clone(ext), &input_refs)?
5644            }
5645            op => tenferro_runtime::extension::apply_standard_op(op.clone(), &input_refs)?,
5646        };
5647        if semantic_outputs.len() != op_node.outputs.len() {
5648            return Err(Error::Internal(format!(
5649                "semantic graph recording expected {} outputs for {:?}, got {}",
5650                op_node.outputs.len(),
5651                op_node.operation,
5652                semantic_outputs.len()
5653            )));
5654        }
5655        for (output_id, output) in op_node.outputs.iter().copied().zip(semantic_outputs) {
5656            values.insert(graph.values()[output_id].key.clone(), output);
5657        }
5658    }
5659
5660    graph
5661        .outputs()
5662        .iter()
5663        .map(|&output_id| {
5664            let key = &graph.values()[output_id].key;
5665            values.get(key).cloned().map(Some).ok_or_else(|| {
5666                Error::Internal(format!(
5667                    "semantic graph recording missing output for {key:?}"
5668                ))
5669            })
5670        })
5671        .collect()
5672}
5673
5674fn record_eager_outputs_from_metadata(
5675    output_metadata: impl IntoIterator<Item = TensorMeta>,
5676    semantic_traces: Vec<Option<TracedTensor>>,
5677    inputs: &[&EagerTensor],
5678) -> Result<RecordedEagerOutputs> {
5679    let output_metadata = output_metadata.into_iter().collect::<Vec<_>>();
5680    if semantic_traces.len() != output_metadata.len() {
5681        return Err(Error::Internal(format!(
5682            "eager recording expected {} semantic traces, got {}",
5683            output_metadata.len(),
5684            semantic_traces.len()
5685        )));
5686    }
5687    let requires_grad =
5688        eager_grad_recording_enabled() && inputs.iter().any(|input| input.requires_grad);
5689    let trace_count = output_metadata.len();
5690    let residual_trace = EagerTrace::new(inputs);
5691    let traces = (0..trace_count)
5692        .map(|_| RecordedEagerTrace {
5693            key: eager_val_key(),
5694            trace: Some(residual_trace.clone()),
5695            requires_grad,
5696        })
5697        .collect();
5698
5699    Ok(RecordedEagerOutputs {
5700        traces,
5701        semantic_traces,
5702    })
5703}
5704
5705fn tensor_meta_from_value(value: &TensorValue) -> TensorMeta {
5706    TensorMeta::exact(
5707        value.dtype(),
5708        value.shape().iter().copied().map(SymDim::from).collect(),
5709    )
5710}
5711
5712#[cfg(test)]
5713pub(crate) fn zero_like_tensor(input: &Tensor, backend: &mut dyn Engine) -> Result<Tensor> {
5714    let host = match input.dtype() {
5715        // A caller-owned payload has no zero-like runtime tensor.
5716        DType::External(type_id) => {
5717            return Err(Error::unsupported(
5718                "zero_like_tensor",
5719                ErrorPhase::GraphBuild,
5720                format!(
5721                    "an externally defined payload ({:?}) has no zero-like runtime tensor",
5722                    DType::External(type_id)
5723                ),
5724            ));
5725        }
5726        DType::F32 => Tensor::from_typed::<f32>(TypedTensor::zeros(input.shape().to_vec())?),
5727        DType::F64 => Tensor::from_typed::<f64>(TypedTensor::zeros(input.shape().to_vec())?),
5728        DType::I32 => Tensor::from_typed::<i32>(TypedTensor::zeros(input.shape().to_vec())?),
5729        DType::I64 => Tensor::from_typed::<i64>(TypedTensor::zeros(input.shape().to_vec())?),
5730        DType::Bool => Tensor::from_typed::<bool>(TypedTensor::from_vec_col_major(
5731            input.shape().to_vec(),
5732            vec![false; input.shape().iter().product()],
5733        )?),
5734        DType::C32 => Tensor::from_typed::<tenferro_tensor::Complex32>(TypedTensor::zeros(
5735            input.shape().to_vec(),
5736        )?),
5737        DType::C64 => Tensor::from_typed::<tenferro_tensor::Complex64>(TypedTensor::zeros(
5738            input.shape().to_vec(),
5739        )?),
5740    };
5741    backend
5742        .upload_host_tensor(TensorRead::from_tensor(&host))
5743        .map_err(Error::from)
5744}
5745
5746pub(crate) fn one_like_tensor(input: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
5747    let host = ones_tensor(input.dtype(), input.shape().to_vec())?;
5748    session
5749        .upload_host_tensor(TensorRead::from_tensor(&host))
5750        .map_err(Error::from)
5751}
5752
5753#[cfg(test)]
5754mod tests;