tenferro_ad/

eager.rs

use std::cell::RefCell;
use std::cmp::Reverse;
use std::collections::HashMap;
use std::env;
use std::fmt;
use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
use std::time::{Duration, Instant};

use crate::extension_cache::{ExtensionCacheLimits, ExtensionCacheStore};
use crate::extension_runtime::{ExtensionExecutor, ExtensionRuntimeRegistryError};
#[cfg(test)]
use computegraph::graph::Graph;
use computegraph::ValueKey;
#[cfg(test)]
use computegraph::ValueRef;
use tenferro_cpu::CpuBackend;
#[cfg(feature = "cuda")]
use tenferro_gpu::CudaBackend;
#[cfg(feature = "webgpu")]
use tenferro_gpu::WebGpuBackend;
use tenferro_ops::input_key::TensorInputKey;
use tenferro_ops::std_tensor_op::StdTensorOp;
use tenferro_ops::ExtensionRuleSet;
use tenferro_ops::ShapeGuardContext;
#[cfg(test)]
use tenferro_tensor::BackendSessionHost;
use tenferro_tensor::{
    CacheStats, DType, Tensor, TensorBackend, TensorElementwise, TensorRead, TensorValue,
    TypedTensor,
};
use tidu::eager::{self, EagerInput, EagerOutput, KeySource, RecordedGraph, Recorder, Trace};

use self::backward::TenferroBackwardCallbacks;
use crate::eager_backend::EagerBackend;
#[cfg(test)]
use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
use crate::eager_exec::{
    exec_op_on_tensor_reads_with_extension_executor, exec_op_on_tensors_with_extension_executor,
};
use crate::error::{ContextId, Error, Result};
#[cfg(test)]
use crate::metadata::push_metadata_scope;
use crate::metadata::{
    metadata_scopes_for_scope, register_scoped_metadata_batch, register_scoped_value_metadata,
    tensor_meta_from_tensor, GlobalMetadataScope,
};
use crate::traced::next_input_key;

use crate::AdContext;

mod backward;

pub(crate) type GradSlot = Arc<Mutex<Option<Arc<Tensor>>>>;
pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<Tensor>>>>;

#[derive(Debug, Default, Clone)]
struct EagerOpProfileEntry {
    calls: usize,
    total_time: Duration,
}

thread_local! {
    static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
        RefCell::new(HashMap::new());
    #[cfg(test)]
    static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
    #[cfg(test)]
    static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
}

pub(crate) fn eager_op_profile_enabled() -> bool {
    #[cfg(test)]
    if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
        return value;
    }

    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
}

pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
    if !eager_op_profile_enabled() {
        return;
    }
    EAGER_OP_PROFILE_STATE.with(|state| {
        let mut state = state.borrow_mut();
        let entry = state.entry(section).or_default();
        entry.calls += 1;
        entry.total_time += elapsed;
    });
}

pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
    if !eager_op_profile_enabled() {
        return f();
    }
    let started = Instant::now();
    let result = f();
    record_eager_op_profile(section, started.elapsed());
    result
}

pub(crate) fn maybe_print_eager_op_profile() {
    if !eager_op_profile_enabled() {
        return;
    }
    let Some(print_every) = eager_op_profile_print_every() else {
        return;
    };
    if print_every == 0 {
        return;
    }

    let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
        state
            .borrow()
            .get("nary_op.total")
            .is_some_and(|entry| entry.calls % print_every == 0)
    });
    if should_print {
        print_and_reset_eager_op_profile();
    }
}

fn eager_op_profile_print_every() -> Option<usize> {
    #[cfg(test)]
    if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
        return value;
    }

    env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
        .ok()?
        .parse()
        .ok()
}

pub(crate) fn print_and_reset_eager_op_profile() {
    EAGER_OP_PROFILE_STATE.with(|state| {
        let mut entries: Vec<_> = state
            .borrow()
            .iter()
            .map(|(section, entry)| (*section, entry.clone()))
            .collect();
        state.borrow_mut().clear();
        entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));

        eprintln!("=== tenferro eager op profile ===");
        for (section, entry) in entries {
            eprintln!(
                "{section}: calls={} total={:.6}ms per_call={:.3}us",
                entry.calls,
                entry.total_time.as_secs_f64() * 1.0e3,
                entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64,
            );
        }
    });
}

/// Stats for caches owned by an [`EagerRuntime`].
///
/// `retained_bytes` fields are logical payload estimates, not process RSS.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct EagerRuntimeCacheStats {
    /// Generic extension runtime caches.
    pub extensions: CacheStats,
}

#[cfg(test)]
pub(crate) struct EagerGraphExecution {
    pub(crate) outputs: Vec<Arc<Tensor>>,
    pub(crate) retained_values: HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
}

/// Shared eager execution context for tensors on a backend.
///
/// Reusing one context lets eager tensors share backend state, extension
/// runtime caches, and gradient storage across a computation.
///
/// # Examples
///
/// ```
/// use tenferro_cpu::CpuBackend;
/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
///
/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
/// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
/// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx).unwrap();
/// let z = x.add(&y).unwrap();
///
/// assert_eq!(z.materialized().unwrap().as_slice::<f64>().unwrap(), &[3.0]);
/// ```
pub struct EagerRuntime {
    pub(crate) backend: Mutex<EagerBackend>,
    pub(crate) extension_executor: Mutex<ExtensionExecutor<EagerBackend>>,
    extension_rules: Option<ExtensionRuleSet>,
    grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
}

impl fmt::Debug for EagerRuntime {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("EagerRuntime");
        match self.backend.try_lock() {
            Ok(backend) => {
                debug.field("backend", &*backend);
            }
            Err(_) => {
                debug.field("backend", &"<locked>");
            }
        }
        match self.extension_executor.try_lock() {
            Ok(executor) => {
                debug.field("extension_executor", &*executor);
            }
            Err(_) => {
                debug.field("extension_executor", &"<locked>");
            }
        }
        debug.field("has_extension_rules", &self.extension_rules.is_some());
        match self.grad_slots.try_lock() {
            Ok(slots) => {
                debug.field("grad_slots_len", &slots.len());
            }
            Err(_) => {
                debug.field("grad_slots_len", &"<locked>");
            }
        }
        debug.finish_non_exhaustive()
    }
}

impl EagerRuntime {
    fn lock_backend(&self) -> Result<MutexGuard<'_, EagerBackend>> {
        self.backend
            .lock()
            .map_err(|_| Error::Internal("backend lock poisoned".to_string()))
    }

    fn lock_extension_executor(&self) -> Result<MutexGuard<'_, ExtensionExecutor<EagerBackend>>> {
        self.extension_executor
            .lock()
            .map_err(|_| Error::Internal("extension executor lock poisoned".to_string()))
    }

    fn lock_grad_slots(
        &self,
    ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
        self.grad_slots
            .lock()
            .map_err(|_| Error::Internal("gradient slot registry lock poisoned".to_string()))
    }

    fn from_backend(backend: EagerBackend) -> Self {
        Self::from_backend_with_extension_rules(backend, None)
    }

    fn from_backend_with_extension_rules(
        backend: EagerBackend,
        extension_rules: Option<ExtensionRuleSet>,
    ) -> Self {
        Self {
            backend: Mutex::new(backend),
            extension_executor: Mutex::new(ExtensionExecutor::new()),
            extension_rules,
            grad_slots: Mutex::new(HashMap::new()),
        }
    }

    /// Create a shared CPU eager execution context.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_ad::EagerRuntime;
    ///
    /// let ctx = EagerRuntime::new();
    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
    /// ```
    pub fn new() -> Arc<Self> {
        Self::with_cpu_backend(CpuBackend::new())
    }

    /// Create a shared eager execution context from a configured CPU backend.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::with_threads(1).unwrap());
    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
    /// ```
    pub fn with_cpu_backend(backend: CpuBackend) -> Arc<Self> {
        Arc::new(Self::from_backend(EagerBackend::cpu(backend)))
    }

    /// Create a shared CPU eager context with explicit AD extension rules.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{AdContext, EagerRuntime};
    ///
    /// let ad = AdContext::builder().build().unwrap();
    /// let ctx = EagerRuntime::with_cpu_backend_and_ad_context(CpuBackend::new(), &ad);
    /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
    /// ```
    pub fn with_cpu_backend_and_ad_context(backend: CpuBackend, ad: &AdContext) -> Arc<Self> {
        Arc::new(Self::from_backend_with_extension_rules(
            EagerBackend::cpu(backend),
            Some(ad.extension_rule_set()),
        ))
    }

    /// Create a shared eager execution context from a configured CUDA backend.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_gpu::CudaBackend;
    /// use tenferro_ad::EagerRuntime;
    ///
    /// let _ctor: fn(CudaBackend) -> std::sync::Arc<EagerRuntime> =
    ///     EagerRuntime::with_cuda_backend;
    /// ```
    #[cfg(feature = "cuda")]
    pub fn with_cuda_backend(backend: CudaBackend) -> Arc<Self> {
        Arc::new(Self::from_backend(EagerBackend::cuda(backend)))
    }

    /// Create a shared CUDA eager context with explicit AD extension rules.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_ad::{AdContext, EagerRuntime};
    /// use tenferro_gpu::CudaBackend;
    ///
    /// let _ctor: fn(CudaBackend, &AdContext) -> std::sync::Arc<EagerRuntime> =
    ///     EagerRuntime::with_cuda_backend_and_ad_context;
    /// ```
    #[cfg(feature = "cuda")]
    pub fn with_cuda_backend_and_ad_context(backend: CudaBackend, ad: &AdContext) -> Arc<Self> {
        Arc::new(Self::from_backend_with_extension_rules(
            EagerBackend::cuda(backend),
            Some(ad.extension_rule_set()),
        ))
    }

    /// Create a shared eager execution context from a configured WebGPU backend.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_ad::EagerRuntime;
    /// use tenferro_gpu::WebGpuBackend;
    ///
    /// let _ctor: fn(WebGpuBackend) -> std::sync::Arc<EagerRuntime> =
    ///     EagerRuntime::with_webgpu_backend;
    /// ```
    #[cfg(feature = "webgpu")]
    pub fn with_webgpu_backend(backend: WebGpuBackend) -> Arc<Self> {
        Arc::new(Self::from_backend(EagerBackend::webgpu(backend)))
    }

    /// Create a shared WebGPU eager context with explicit AD extension rules.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_ad::{AdContext, EagerRuntime};
    /// use tenferro_gpu::WebGpuBackend;
    ///
    /// let _ctor: fn(WebGpuBackend, &AdContext) -> std::sync::Arc<EagerRuntime> =
    ///     EagerRuntime::with_webgpu_backend_and_ad_context;
    /// ```
    #[cfg(feature = "webgpu")]
    pub fn with_webgpu_backend_and_ad_context(backend: WebGpuBackend, ad: &AdContext) -> Arc<Self> {
        Arc::new(Self::from_backend_with_extension_rules(
            EagerBackend::webgpu(backend),
            Some(ad.extension_rule_set()),
        ))
    }

    /// Return an opaque identifier for this context.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// assert_ne!(ctx.id(), EagerRuntime::with_cpu_backend(CpuBackend::new()).id());
    /// ```
    pub fn id(&self) -> ContextId {
        ContextId::from_ptr(self)
    }

    /// Register one extension runtime on this eager context.
    pub fn register_extension(
        &self,
        register: impl FnOnce(
            &mut ExtensionExecutor<EagerBackend>,
        ) -> std::result::Result<(), ExtensionRuntimeRegistryError>,
    ) -> std::result::Result<(), ExtensionRuntimeRegistryError> {
        let mut executor = self.extension_executor.lock().map_err(|_| {
            ExtensionRuntimeRegistryError::PoisonedLock {
                name: "extension executor lock",
            }
        })?;
        register(&mut executor)
    }

    /// Clear generic extension runtime cache entries.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// ctx.clear_extension_caches()?;
    /// assert_eq!(ctx.cache_stats()?.extensions.entries, 0);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn clear_extension_caches(&self) -> Result<()> {
        self.lock_extension_executor()?.clear_caches();
        Ok(())
    }

    /// Clear every cache owned by this eager context.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// ctx.clear_caches()?;
    /// assert_eq!(ctx.cache_stats()?.extensions.entries, 0);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn clear_caches(&self) -> Result<()> {
        self.clear_extension_caches()
    }

    /// Return eager runtime cache-entry and retained-byte stats.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let stats = ctx.cache_stats()?;
    /// assert_eq!(stats.extensions.entries, 0);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
        Ok(EagerRuntimeCacheStats {
            extensions: self.lock_extension_executor()?.cache_stats(),
        })
    }

    /// Return the extension cache retention limits.
    pub fn extension_cache_limits(&self) -> Result<ExtensionCacheLimits> {
        Ok(self.lock_extension_executor()?.cache_limits())
    }

    /// Replace extension cache retention limits.
    pub fn set_extension_cache_limits(&self, limits: ExtensionCacheLimits) -> Result<()> {
        self.lock_extension_executor()?.set_cache_limits(limits);
        Ok(())
    }

    /// Mutably borrow generic extension runtime cache storage.
    ///
    /// This hook is for standard extension crates that need cache entries
    /// owned by an eager runtime while preserving eager value semantics outside
    /// a registered extension execution boundary.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_ad::EagerRuntime;
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_runtime::ExtensionCacheKey;
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let key = ExtensionCacheKey::new("example.cache.v1", "plans", 1);
    ///
    /// ctx.with_extension_caches_mut(|caches| {
    ///     caches.put(key, 7_usize, std::mem::size_of::<usize>());
    /// });
    ///
    /// assert_eq!(ctx.cache_stats().unwrap().extensions.entries, 1);
    /// ```
    pub fn with_extension_caches_mut<R>(
        &self,
        f: impl FnOnce(&mut ExtensionCacheStore) -> R,
    ) -> Result<R> {
        let mut executor = self.lock_extension_executor()?;
        Ok(f(executor.caches_mut()))
    }

    /// Mutably borrow this runtime's backend.
    ///
    /// This hook lets standard extension crates run a whole contraction program
    /// in a single backend session (instead of one eager op per step) while
    /// preserving eager value semantics for untracked tensors.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_ad::EagerRuntime;
    /// use tenferro_cpu::CpuBackend;
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// // The closure receives `&mut EagerBackend`; standard extension crates
    /// // use it to open one backend session for a whole contraction program.
    /// let answer = ctx.with_backend_mut(|_backend| 42).unwrap();
    /// assert_eq!(answer, 42);
    /// ```
    pub fn with_backend_mut<R>(&self, f: impl FnOnce(&mut EagerBackend) -> R) -> Result<R> {
        let mut backend = self.lock_backend()?;
        Ok(f(&mut backend))
    }

    /// Block the current thread until backend work submitted by this eager runtime completes.
    ///
    /// CPU runtimes return immediately. CUDA and WebGPU runtimes synchronize
    /// their current backend work queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::EagerRuntime;
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// ctx.synchronize().unwrap();
    /// ```
    pub fn synchronize(&self) -> Result<()> {
        self.lock_backend()?.synchronize().map_err(Error::from)
    }

    pub(crate) fn exec_outputs(&self, op: &StdTensorOp, inputs: &[&Tensor]) -> Result<Vec<Tensor>> {
        let mut backend =
            profile_eager_op_section("exec_outputs.lock_backend", || self.lock_backend())?;
        let mut extension_executor =
            profile_eager_op_section("exec_outputs.lock_extensions", || {
                self.lock_extension_executor()
            })?;
        profile_eager_op_section("exec_outputs.exec_op", || {
            exec_op_on_tensors_with_extension_executor(
                op,
                inputs,
                &mut *backend,
                Some(&mut *extension_executor),
            )
        })
    }

    pub(crate) fn exec_outputs_read(
        &self,
        op: &StdTensorOp,
        inputs: &[TensorRead<'_>],
    ) -> Result<Vec<Tensor>> {
        let mut backend =
            profile_eager_op_section("exec_outputs_read.lock_backend", || self.lock_backend())?;
        let mut extension_executor =
            profile_eager_op_section("exec_outputs_read.lock_extensions", || {
                self.lock_extension_executor()
            })?;
        profile_eager_op_section("exec_outputs_read.exec_op", || {
            exec_op_on_tensor_reads_with_extension_executor(
                op,
                inputs,
                &mut *backend,
                Some(&mut *extension_executor),
            )
        })
    }

    #[cfg(test)]
    pub(crate) fn exec_standard_graph_outputs(
        &self,
        graph: &Graph<StdTensorOp>,
        initial_data: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
    ) -> Result<EagerGraphExecution> {
        let mut backend =
            profile_eager_op_section("exec_graph.lock_backend", || self.lock_backend())?;
        let mut all_values = initial_data.clone();

        profile_eager_op_section("exec_graph.with_backend_session", || {
            backend.with_backend_session(|exec| -> Result<()> {
                for op_node in graph.operations() {
                    let outputs = {
                        let input_values = op_node
                            .inputs
                            .iter()
                            .map(|input| {
                                let key = match input {
                                    ValueRef::Local(local_id) => &graph.values()[*local_id].key,
                                    ValueRef::External(key) => key,
                                };
                                all_values.get(key).cloned().ok_or_else(|| {
                                    Error::Internal(format!(
                                        "standard graph eager execution missing value for {key:?}"
                                    ))
                                })
                            })
                            .collect::<Result<Vec<_>>>()?;
                        let input_reads = input_values
                            .iter()
                            .map(|value| TensorRead::from_tensor(value.as_ref()))
                            .collect::<Vec<_>>();
                        exec_standard_op_on_tensor_reads_in_session(
                            &op_node.operation,
                            &input_reads,
                            exec,
                        )?
                    };

                    if outputs.len() != op_node.outputs.len() {
                        return Err(Error::Internal(format!(
                            "standard graph eager execution expected {} outputs for {:?}, got {}",
                            op_node.outputs.len(),
                            op_node.operation,
                            outputs.len()
                        )));
                    }

                    for (output_id, output) in op_node.outputs.iter().zip(outputs) {
                        let key = graph.values()[*output_id].key.clone();
                        all_values.insert(key, Arc::new(output));
                    }
                }
                Ok(())
            })
        })?;

        let outputs = graph
            .outputs()
            .iter()
            .map(|&output_id| {
                let key = &graph.values()[output_id].key;
                all_values.get(key).cloned().ok_or_else(|| {
                    Error::Internal(format!(
                        "standard graph eager execution missing graph output {key:?}"
                    ))
                })
            })
            .collect::<Result<Vec<_>>>()?;

        Ok(EagerGraphExecution {
            outputs,
            retained_values: all_values,
        })
    }

    pub(crate) fn try_register_grad_slot(
        &self,
        key: &ValueKey<StdTensorOp>,
        slot: &GradSlot,
    ) -> Result<()> {
        self.lock_grad_slots()?
            .insert(key.clone(), Arc::downgrade(slot));
        Ok(())
    }

    /// Clear all live gradient slots tracked by this context.
    ///
    /// This resets the stored gradients to `None` without unregistering the
    /// tensors, so future `backward()` calls can accumulate again.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// 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();
    /// 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();
    /// let loss = x.mul(&y).unwrap().reduce_sum(&[0]).unwrap();
    /// let _ = loss.backward().unwrap();
    ///
    /// ctx.clear_grads()?;
    ///
    /// assert!(x.grad()?.is_none());
    /// assert!(y.grad()?.is_none());
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn clear_grads(&self) -> Result<()> {
        let mut poisoned_slot = false;
        self.lock_grad_slots()?.retain(|_, slot| {
            if let Some(slot) = slot.upgrade() {
                match slot.lock() {
                    Ok(mut current) => {
                        *current = None;
                    }
                    Err(_) => {
                        poisoned_slot = true;
                    }
                }
                true
            } else {
                false
            }
        });
        if poisoned_slot {
            return Err(Error::Internal("gradient slot lock poisoned".to_string()));
        }
        Ok(())
    }

    /// Import a concrete tensor into this context as an untracked constant.
    ///
    /// The returned tensor does not participate in gradient tracking.
    /// Use this for fixed masks, quadrature weights, physical constants,
    /// and other data that should not receive gradients.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let c = ctx.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx)?;
    /// let z = x.add(&c).unwrap();
    ///
    /// assert_eq!(z.materialized()?.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
        EagerTensor::new_leaf(Arc::clone(self), tensor, false)
    }

    /// Import a concrete tensor into this context as a trainable variable.
    ///
    /// The returned tensor participates in gradient tracking; its gradient
    /// slot is registered in this context.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let p = ctx.variable_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
    /// let loss = p.exp().unwrap().reduce_sum(&[0]).unwrap();
    /// let _ = loss.backward().unwrap();
    ///
    /// let grad = p.grad().unwrap().unwrap();
    /// assert_eq!(grad.shape(), &[2]);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
        EagerTensor::new_leaf(Arc::clone(self), tensor, true)
    }

    fn store_grads(
        &self,
        cotangents: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
        backend: &mut EagerBackend,
    ) -> Result<()> {
        let mut updates = Vec::new();

        {
            let mut slots = self.lock_grad_slots()?;
            slots.retain(|key, slot| {
                let Some(slot) = slot.upgrade() else {
                    return false;
                };

                if let Some(incoming) = cotangents.get(key) {
                    updates.push((slot, Arc::clone(incoming)));
                }

                true
            });
        }

        for (slot, incoming) in updates {
            let mut current = slot
                .lock()
                .map_err(|_| Error::Internal("gradient slot lock poisoned".to_string()))?;
            let next = match current.as_ref() {
                Some(existing) => Arc::new(backend.add(existing.as_ref(), incoming.as_ref())?),
                None => incoming,
            };
            *current = Some(next);
        }

        Ok(())
    }
}

/// Eager tensor with reverse-mode autodiff over concrete tensor values.
///
/// This executes each primitive immediately and records a lightweight reverse
/// DAG for `backward()`. Gradients accumulate across repeated `backward()`
/// calls until they are cleared explicitly.
///
/// # Examples
///
/// ```
/// use tenferro_cpu::CpuBackend;
/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
///
/// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
/// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx)?;
/// let loss = x.mul(&x).unwrap().reduce_sum(&[0]).unwrap();
/// let _cotangents = loss.backward().unwrap();
/// let loss = x.mul(&x).unwrap().reduce_sum(&[0]).unwrap();
/// let _cotangents = loss.backward().unwrap();
///
/// assert_eq!(x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(), &[4.0, 8.0, 12.0]);
/// x.clear_grad();
///
/// assert!(x.grad().unwrap().is_none());
/// # Ok::<(), tenferro_ad::Error>(())
/// ```
#[derive(Clone)]
pub struct EagerTensor {
    pub(crate) value: Arc<TensorValue>,
    materialized_cache: Arc<OnceLock<Arc<Tensor>>>,
    pub(crate) key: ValueKey<StdTensorOp>,
    pub(crate) trace: Option<Trace<StdTensorOp>>,
    pub(crate) requires_grad: bool,
    grad_slot: GradSlot,
    pub(crate) metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
    pub(crate) ctx: Arc<EagerRuntime>,
}

impl fmt::Debug for EagerTensor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("EagerTensor")
            .field("dtype", &self.dtype())
            .field("shape", &self.shape())
            .field("key", &self.key)
            .field("requires_grad", &self.requires_grad)
            .field("has_trace", &self.trace.is_some())
            .field("ctx_id", &self.ctx_id())
            .finish_non_exhaustive()
    }
}

impl EagerTensor {
    /// Create an untracked eager tensor inside an existing eager context.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
    ///
    /// assert_eq!(x.materialized()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
        Self::new_leaf(ctx, tensor, false)
    }

    /// Create a tracked eager leaf inside an existing eager context.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
    ///
    /// assert!(x.grad().unwrap().is_none());
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
        Self::new_leaf(ctx, tensor, true)
    }

    pub(crate) fn new_leaf(
        ctx: Arc<EagerRuntime>,
        tensor: Tensor,
        requires_grad: bool,
    ) -> Result<Self> {
        let key = eager_val_key();
        let metadata_scope =
            register_scoped_value_metadata(key.clone(), tensor_meta_from_tensor(&tensor)).map_err(
                |err| Error::Internal(format!("eager leaf metadata registration failed: {err}")),
            )?;
        let tensor = Arc::new(tensor);
        let grad_slot = Arc::new(Mutex::new(None));
        if requires_grad {
            ctx.try_register_grad_slot(&key, &grad_slot)?;
        }

        Ok(Self {
            value: Arc::new(TensorValue::from_tensor_arc(tensor)),
            materialized_cache: Arc::new(OnceLock::new()),
            key,
            trace: None,
            requires_grad,
            grad_slot,
            metadata_scopes: metadata_scopes_for_scope(metadata_scope),
            ctx,
        })
    }

    pub(crate) fn new_result(
        ctx: Arc<EagerRuntime>,
        key: ValueKey<StdTensorOp>,
        tensor: Tensor,
        requires_grad: bool,
        trace: Option<Trace<StdTensorOp>>,
        metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
    ) -> Result<Self> {
        Self::new_result_arc(
            ctx,
            key,
            Arc::new(tensor),
            requires_grad,
            trace,
            metadata_scopes,
        )
    }

    pub(crate) fn new_result_arc(
        ctx: Arc<EagerRuntime>,
        key: ValueKey<StdTensorOp>,
        tensor: Arc<Tensor>,
        requires_grad: bool,
        trace: Option<Trace<StdTensorOp>>,
        metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
    ) -> Result<Self> {
        let grad_slot = Arc::new(Mutex::new(None));
        if requires_grad {
            ctx.try_register_grad_slot(&key, &grad_slot)?;
        }

        Ok(Self {
            value: Arc::new(TensorValue::from_tensor_arc(tensor)),
            materialized_cache: Arc::new(OnceLock::new()),
            key,
            trace,
            requires_grad,
            grad_slot,
            metadata_scopes,
            ctx,
        })
    }

    pub(crate) fn new_result_value(
        ctx: Arc<EagerRuntime>,
        key: ValueKey<StdTensorOp>,
        value: TensorValue,
        requires_grad: bool,
        trace: Option<Trace<StdTensorOp>>,
        metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
    ) -> Result<Self> {
        let grad_slot = Arc::new(Mutex::new(None));
        if requires_grad {
            ctx.try_register_grad_slot(&key, &grad_slot)?;
        }

        Ok(Self {
            value: Arc::new(value),
            materialized_cache: Arc::new(OnceLock::new()),
            key,
            trace,
            requires_grad,
            grad_slot,
            metadata_scopes,
            ctx,
        })
    }

    pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
        Self::new_result(ctx, eager_val_key(), tensor, false, None, Vec::new())
    }

    pub(crate) fn new_untracked_value_result(ctx: Arc<EagerRuntime>, value: TensorValue) -> Self {
        Self {
            value: Arc::new(value),
            materialized_cache: Arc::new(OnceLock::new()),
            key: eager_val_key(),
            trace: None,
            requires_grad: false,
            grad_slot: Arc::new(Mutex::new(None)),
            metadata_scopes: Vec::new(),
            ctx,
        }
    }

    /// Detach this tensor from the reverse graph.
    ///
    /// The returned tensor keeps the concrete value but no longer contributes
    /// gradients to the original graph.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
    /// let y = x.detach();
    ///
    /// assert_eq!(y.materialized()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
    /// assert!(y.grad().unwrap().is_none());
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn detach(&self) -> Self {
        Self::new_untracked_value_result(self.ctx.clone(), self.value.as_ref().clone())
    }

    /// Detach this tensor from its graph and re-register it in a different
    /// context as an untracked leaf.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx_a = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let ctx_b = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx_a)?;
    /// let d = x.detach_into(&ctx_b)?;
    ///
    /// assert!(!d.tracks_grad());
    /// assert_eq!(d.ctx_id(), ctx_b.id());
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn detach_into(&self, ctx: &Arc<EagerRuntime>) -> Result<Self> {
        Self::from_tensor_in(self.to_tensor()?, Arc::clone(ctx))
    }

    /// Materialize and share the concrete tensor value.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap(), ctx)?;
    /// assert_eq!(x.materialized()?.as_slice::<f64>().unwrap(), &[3.0]);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn materialized(&self) -> Result<Arc<Tensor>> {
        self.materialized_arc()
    }

    /// Return this tensor's scalar dtype without materializing through
    /// [`materialized`](Self::materialized).
    pub fn dtype(&self) -> DType {
        self.value.dtype()
    }

    /// Return this tensor's logical shape without materializing through
    /// [`materialized`](Self::materialized).
    pub fn shape(&self) -> &[usize] {
        self.value.shape()
    }

    /// Borrow this tensor value as a [`TensorRead`].
    ///
    /// This is the preferred borrowed input boundary for executor calls. It
    /// preserves the option to replace eager storage with non-contiguous views
    /// without forcing callers through [`materialized`](Self::materialized).
    pub fn tensor_read(&self) -> TensorRead<'_> {
        self.value.tensor_read()
    }

    /// Materialize this eager tensor as an owned [`Tensor`].
    ///
    /// This is the owned materialization boundary for callers that need a
    /// standalone compact tensor. The operation is fallible because eager
    /// values may be backed by lazy or backend-resident storage.
    pub fn to_tensor(&self) -> Result<Tensor> {
        self.value.to_tensor().map_err(Error::from)
    }

    pub(crate) fn materialized_arc(&self) -> Result<Arc<Tensor>> {
        if let Some(tensor) = self.value.as_tensor_arc() {
            return Ok(Arc::clone(tensor));
        }
        if let Some(tensor) = self.materialized_cache.get() {
            return Ok(Arc::clone(tensor));
        }

        let materialized = Arc::new(self.value.to_tensor().map_err(Error::from)?);
        let _ = self.materialized_cache.set(Arc::clone(&materialized));
        Ok(self
            .materialized_cache
            .get()
            .map(Arc::clone)
            .unwrap_or(materialized))
    }

    #[cfg(test)]
    pub(crate) fn materialized_cache_is_initialized(&self) -> bool {
        self.materialized_cache.get().is_some()
    }

    /// Return the accumulated gradient currently stored for this tensor.
    ///
    /// The stored gradient accumulates across repeated `backward()` calls
    /// until it is cleared explicitly.
    ///
    /// For complex scalar losses, stored gradients use tenferro's
    /// Hermitian-adjoint cotangent convention. See
    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx).unwrap();
    /// let loss = x.exp().unwrap().reduce_sum(&[0]).unwrap();
    /// let _cotangents = loss.backward().unwrap();
    ///
    /// let grad = x.grad()?.unwrap();
    /// assert_eq!(grad.shape(), &[2]);
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn grad(&self) -> Result<Option<Arc<Tensor>>> {
        self.grad_slot
            .lock()
            .map_err(|_| Error::Internal("gradient slot lock poisoned".to_string()))
            .map(|slot| slot.clone())
    }

    /// Clear the accumulated gradient stored for this tensor.
    ///
    /// This only affects this tensor's gradient slot. Other tensors in the
    /// same context retain their gradients until they are cleared explicitly or
    /// overwritten by later accumulation.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// 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();
    /// let y = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![4.0_f64, 5.0, 6.0]).unwrap(), ctx).unwrap();
    /// let loss = x.mul(&y).unwrap().reduce_sum(&[0]).unwrap();
    /// let _ = loss.backward().unwrap();
    ///
    /// x.clear_grad()?;
    ///
    /// assert!(x.grad()?.is_none());
    /// assert!(y.grad()?.is_some());
    /// # Ok::<(), tenferro_ad::Error>(())
    /// ```
    pub fn clear_grad(&self) -> Result<()> {
        *self
            .grad_slot
            .lock()
            .map_err(|_| Error::Internal("gradient slot lock poisoned".to_string()))? = None;
        Ok(())
    }

    /// Report whether this tensor participates in gradient tracking.
    ///
    /// Tracked tensors keep a gradient slot in their eager context; untracked
    /// tensors and detached tensors do not.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let plain = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
    /// let tracked = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
    /// let detached = tracked.detach();
    ///
    /// assert!(!plain.tracks_grad());
    /// assert!(tracked.tracks_grad());
    /// assert!(!detached.tracks_grad());
    /// ```
    pub fn tracks_grad(&self) -> bool {
        self.requires_grad
    }

    #[cfg(test)]
    fn debug_trace_saved_value_count(&self) -> Option<usize> {
        self.trace.as_ref().map(|trace| trace.saved_values().len())
    }

    /// Return the opaque identifier of the context this tensor belongs to.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
    ///
    /// assert_eq!(x.ctx_id(), ctx.id());
    /// ```
    pub fn ctx_id(&self) -> ContextId {
        self.ctx.id()
    }

    /// Borrow the eager runtime context that owns this tensor.
    pub fn runtime(&self) -> &Arc<EagerRuntime> {
        &self.ctx
    }

    /// Check whether two tensors belong to the same eager context.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
    /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx).unwrap();
    ///
    /// assert!(x.same_context(&y));
    /// ```
    pub fn same_context(&self, other: &Self) -> bool {
        self.ctx_id() == other.ctx_id()
    }

    #[cfg(test)]
    pub(crate) fn standard_graph_op(
        inputs: &[&Self],
        build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
    ) -> Result<Vec<Self>> {
        let Some(first) = inputs.first() else {
            return Err(Error::Internal(
                "standard eager graph op requires at least one input tensor".to_string(),
            ));
        };
        let ctx = Arc::clone(&first.ctx);
        for tensor in inputs.iter().skip(1) {
            if !first.same_context(tensor) {
                return Err(Error::ContextMismatch {
                    lhs: first.ctx_id(),
                    rhs: tensor.ctx_id(),
                });
            }
        }

        let mut recorder = Recorder::new(EagerTensorKeySource);
        let graph_input_keys = recorder.fresh_input_keys::<StdTensorOp>(inputs.len());
        let graph = build_graph(&graph_input_keys)?;
        let initial_data = graph_input_keys
            .iter()
            .zip(inputs.iter())
            .map(|(key, tensor)| Ok((ValueKey::Input(key.clone()), tensor.materialized_arc()?)))
            .collect::<Result<HashMap<_, _>>>()?;
        let execution = ctx.exec_standard_graph_outputs(graph.as_ref(), &initial_data)?;
        if execution.outputs.len() != graph.outputs().len() {
            return Err(Error::Internal(format!(
                "standard eager graph op expected {} graph outputs, got {}",
                graph.outputs().len(),
                execution.outputs.len()
            )));
        }

        if !inputs.iter().any(|input| input.requires_grad) {
            return execution
                .outputs
                .into_iter()
                .map(|output| {
                    Self::new_result_arc(
                        Arc::clone(&ctx),
                        eager_val_key(),
                        output,
                        false,
                        None,
                        Vec::new(),
                    )
                })
                .collect();
        }

        let output_keys = graph
            .outputs()
            .iter()
            .map(|&output_id| graph.values()[output_id].key.clone())
            .collect();
        let recorded_graph = RecordedGraph::new(Arc::clone(&graph), graph_input_keys, output_keys)
            .map_err(eager_record_error)?;
        let recorded = record_eager_recorded_graph_outputs(
            &mut recorder,
            recorded_graph,
            &execution.outputs,
            execution.retained_values,
            inputs,
        )?;
        if recorded.traces.len() != execution.outputs.len() {
            return Err(Error::Internal(format!(
                "standard eager graph op expected {} eager traces, got {}",
                execution.outputs.len(),
                recorded.traces.len()
            )));
        }

        let mut metadata_scopes = vec![Arc::clone(&recorded.metadata_scope)];
        for input in inputs {
            for scope in &input.metadata_scopes {
                push_metadata_scope(&mut metadata_scopes, Arc::clone(scope));
            }
        }

        recorded
            .traces
            .into_iter()
            .zip(execution.outputs)
            .map(|(trace, output)| {
                Self::new_result_arc(
                    Arc::clone(&ctx),
                    trace.key,
                    output,
                    trace.requires_grad,
                    trace.trace,
                    metadata_scopes.clone(),
                )
            })
            .collect()
    }

    /// Run reverse-mode AD from this scalar output.
    ///
    /// Returns the full cotangent map produced by the reverse pass and also
    /// accumulates into `grad()` for tracked eager tensors reachable from this
    /// output.
    ///
    /// For complex scalar outputs, cotangents use tenferro's Hermitian
    /// real-inner-product convention. See
    /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_cpu::CpuBackend;
    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
    ///
    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
    /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx).unwrap();
    /// let loss = x.add(&x).unwrap().reduce_sum(&[0]).unwrap();
    /// let _cotangents = loss.backward().unwrap();
    /// let loss = x.add(&x).unwrap().reduce_sum(&[0]).unwrap();
    /// let _cotangents = loss.backward().unwrap();
    ///
    /// assert_eq!(x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(), &[4.0, 4.0, 4.0]);
    /// ```
    pub fn backward(&self) -> Result<HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>> {
        if !self.shape().is_empty() {
            return Err(Error::NonScalarGrad {
                shape: self.shape().to_vec(),
            });
        }

        let value = self.materialized_arc()?;
        let mut backend = self.ctx.lock_backend()?;
        let mut extension_executor = self.ctx.lock_extension_executor()?;
        let seed = Arc::new(one_like_tensor(value.as_ref(), &mut *backend)?);
        let mut callbacks = TenferroBackwardCallbacks::new(
            &mut *backend,
            Some(&mut *extension_executor),
            self.metadata_scopes.clone(),
        );
        let mut ad_ctx = ShapeGuardContext::with_global_metadata();
        if let Some(extension_rules) = &self.ctx.extension_rules {
            ad_ctx = ad_ctx.with_extension_rules(extension_rules.clone());
        }
        let cotangents_result = eager::backward(
            &self.key,
            self.trace.as_ref(),
            seed,
            &mut callbacks,
            &mut ad_ctx,
        );
        let callback_error = callbacks.take_error();
        drop(callbacks);
        let cotangents = match (cotangents_result, callback_error) {
            (_, Some(err)) => return Err(Error::Internal(err.to_string())),
            (Err(err), None) => return Err(Error::Internal(err.to_string())),
            (Ok(cotangents), None) => cotangents,
        };
        self.ctx.store_grads(&cotangents, &mut backend)?;
        Ok(cotangents)
    }
}

pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
    ValueKey::Input(next_input_key())
}

pub(crate) struct EagerTensorKeySource;

impl KeySource<StdTensorOp> for EagerTensorKeySource {
    fn fresh_input_key(&mut self) -> TensorInputKey {
        next_input_key()
    }
}

pub(crate) fn eager_value(tensor: &EagerTensor) -> Result<EagerInput<StdTensorOp>> {
    Ok(EagerInput {
        key: tensor.key.clone(),
        trace: tensor.trace.clone(),
        requires_grad: tensor.requires_grad,
        data: tensor.materialized_arc()?,
    })
}

pub(crate) struct RecordedEagerOutputs {
    pub(crate) traces: Vec<EagerOutput<StdTensorOp>>,
    pub(crate) metadata_scope: Arc<GlobalMetadataScope>,
}

pub(crate) fn record_eager_outputs(
    op: &StdTensorOp,
    outputs: &[Arc<Tensor>],
    inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
    let mut recorder = Recorder::new(EagerTensorKeySource);
    let graph_input_keys = recorder.fresh_input_keys::<StdTensorOp>(inputs.len());
    let graph =
        RecordedGraph::from_primitive(op.clone(), graph_input_keys).map_err(eager_record_error)?;
    let retained_values = graph
        .output_keys()
        .iter()
        .cloned()
        .zip(outputs.iter().cloned())
        .collect();
    record_eager_recorded_graph_outputs(&mut recorder, graph, outputs, retained_values, inputs)
}

pub(crate) fn record_eager_recorded_graph_outputs(
    recorder: &mut Recorder<EagerTensorKeySource>,
    graph: RecordedGraph<StdTensorOp>,
    outputs: &[Arc<Tensor>],
    retained_values: HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
    inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
    let input_values: Vec<_> = inputs
        .iter()
        .map(|tensor| eager_value(tensor))
        .collect::<Result<_>>()?;
    let traces = recorder
        .record_graph(graph, &input_values, outputs, retained_values)
        .map_err(eager_record_error)?;

    let mut registrations = Vec::new();
    for trace in &traces {
        if let Some(output) = outputs.get(trace.output_slot) {
            registrations.push((trace.key.clone(), tensor_meta_from_tensor(output.as_ref())));
        }
    }

    if let Some(trace) = traces.iter().find_map(|output| output.trace.as_ref()) {
        for (key, value) in trace.saved_values() {
            registrations.push((key.clone(), tensor_meta_from_tensor(value.as_ref())));
        }
    }

    Ok(RecordedEagerOutputs {
        traces,
        metadata_scope: Arc::new(register_scoped_metadata_batch(registrations)?),
    })
}

fn eager_record_error(err: tidu::eager::EagerRecordError) -> Error {
    Error::Internal(format!("invalid eager recording metadata: {err}"))
}

pub(crate) fn exec_single_output(
    op: &StdTensorOp,
    inputs: &[&Tensor],
    ctx: &EagerRuntime,
) -> Result<Tensor> {
    let mut outputs = ctx.exec_outputs(op, inputs)?;
    if outputs.len() != 1 {
        return Err(Error::Internal(format!(
            "expected one eager output for {:?}, got {}",
            op,
            outputs.len()
        )));
    }
    Ok(profile_eager_op_section(
        "exec_single_output.remove_output",
        || outputs.remove(0),
    ))
}

pub(crate) fn exec_single_output_read(
    op: &StdTensorOp,
    inputs: &[TensorRead<'_>],
    ctx: &EagerRuntime,
) -> Result<Tensor> {
    let mut outputs = ctx.exec_outputs_read(op, inputs)?;
    if outputs.len() != 1 {
        return Err(Error::Internal(format!(
            "expected one eager output for {:?}, got {}",
            op,
            outputs.len()
        )));
    }
    Ok(profile_eager_op_section(
        "exec_single_output_read.remove_output",
        || outputs.remove(0),
    ))
}

pub(crate) fn zero_like_tensor<B: TensorBackend>(
    input: &Tensor,
    backend: &mut B,
) -> Result<Tensor> {
    let host = match input {
        Tensor::F32(tensor) => Tensor::F32(TypedTensor::zeros(tensor.shape().to_vec())?),
        Tensor::F64(tensor) => Tensor::F64(TypedTensor::zeros(tensor.shape().to_vec())?),
        Tensor::I32(tensor) => Tensor::I32(TypedTensor::zeros(tensor.shape().to_vec())?),
        Tensor::I64(tensor) => Tensor::I64(TypedTensor::zeros(tensor.shape().to_vec())?),
        Tensor::Bool(tensor) => Tensor::Bool(TypedTensor::from_vec_col_major(
            tensor.shape().to_vec(),
            vec![false; tensor.n_elements()],
        )?),
        Tensor::C32(tensor) => Tensor::C32(TypedTensor::zeros(tensor.shape().to_vec())?),
        Tensor::C64(tensor) => Tensor::C64(TypedTensor::zeros(tensor.shape().to_vec())?),
    };
    backend.upload_host_tensor(&host).map_err(Error::from)
}

pub(crate) fn one_like_tensor<B: TensorBackend>(input: &Tensor, backend: &mut B) -> Result<Tensor> {
    let zero = zero_like_tensor(input, backend)?;
    backend.exp(&zero).map_err(Error::from)
}

#[cfg(test)]
mod tests;