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