tenferro/

traced.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use computegraph::compile::compile;
use computegraph::fragment::{Fragment, FragmentBuilder};
use computegraph::materialize::materialize_merge;
use computegraph::resolve::resolve;
use computegraph::types::{GlobalValKey, OpMode, ValRef};
use computegraph::LocalValId;
use num_complex::{Complex32, Complex64};
use tenferro_ops::dim_expr::DimExpr;
use tenferro_ops::input_key::TensorInputKey;
use tenferro_ops::std_tensor_op::StdTensorOp;
use tenferro_ops::ShapeGuardContext;
use tenferro_tensor::{DType, DotGeneralConfig, Tensor, TensorBackend, TensorScalar, TypedTensor};
use tidu::{differentiate, transpose};

use super::compiler::compile_std_to_exec;
use super::engine::Engine;
use super::error::{Error, Result};
use super::sym_dim::SymDim;
use crate::checkpoint::CheckpointNode;

static NEXT_INPUT_ID: AtomicU64 = AtomicU64::new(0);
static NEXT_DIFF_PASS_ID: AtomicU64 = AtomicU64::new(0);
static NEXT_TRACED_ID: AtomicU64 = AtomicU64::new(0);

pub type TracedTensorId = u64;

pub(crate) fn next_input_key() -> TensorInputKey {
    TensorInputKey::User {
        id: NEXT_INPUT_ID.fetch_add(1, Ordering::Relaxed),
    }
}

fn next_pass_id() -> u64 {
    NEXT_DIFF_PASS_ID.fetch_add(1, Ordering::Relaxed)
}

pub(crate) fn next_traced_id() -> TracedTensorId {
    NEXT_TRACED_ID.fetch_add(1, Ordering::Relaxed)
}

#[derive(Clone)]
pub struct TracedTensor {
    pub id: TracedTensorId,
    pub rank: usize,
    pub dtype: DType,
    pub fragment: Arc<Fragment<StdTensorOp>>,
    pub val: LocalValId,
    pub data: Option<Arc<Tensor>>,
    pub(crate) shape_hint: Option<Vec<SymDim>>,
    pub(crate) inputs_map: Arc<HashMap<TensorInputKey, Arc<Tensor>>>,
    pub(crate) extra_roots: Vec<Arc<Fragment<StdTensorOp>>>,
    pub(crate) checkpoint_chain: Option<Arc<CheckpointNode>>,
}

/// Compute a broadcast output shape following NumPy rules.
///
/// Returns `None` when the two shapes are incompatible.
fn broadcast_shape(a: &[usize], b: &[usize]) -> Option<Vec<usize>> {
    let rank = a.len().max(b.len());
    let mut result = Vec::with_capacity(rank);
    for index in 0..rank {
        let a_dim = if index < rank - a.len() {
            1
        } else {
            a[index - (rank - a.len())]
        };
        let b_dim = if index < rank - b.len() {
            1
        } else {
            b[index - (rank - b.len())]
        };
        if a_dim == b_dim {
            result.push(a_dim);
        } else if a_dim == 1 {
            result.push(b_dim);
        } else if b_dim == 1 {
            result.push(a_dim);
        } else {
            return None;
        }
    }
    Some(result)
}

pub(crate) fn try_concrete_shape(tensor: &TracedTensor) -> Option<Vec<usize>> {
    tensor
        .shape_hint
        .as_ref()?
        .iter()
        .map(SymDim::constant_value)
        .collect()
}

pub(crate) fn concrete_shape(tensor: &TracedTensor) -> Vec<usize> {
    tensor
        .shape_hint
        .as_ref()
        .unwrap_or_else(|| panic!("missing shape hint for traced tensor {}", tensor.id))
        .iter()
        .map(|dim| {
            dim.constant_value().unwrap_or_else(|| {
                panic!("symbolic dimension in shape hint for tensor {}", tensor.id)
            })
        })
        .collect()
}

fn error_shape_hint(tensor: &TracedTensor) -> Vec<usize> {
    try_concrete_shape(tensor).unwrap_or_else(|| vec![0; tensor.rank])
}

/// Broadcast a traced tensor to `target_shape`.
///
/// Expanding singleton axes are first reshaped away so the existing
/// `BroadcastInDim` transpose rule reduces them correctly during VJP.
fn broadcast_to(tensor: &TracedTensor, target_shape: &[usize]) -> TracedTensor {
    let tensor_shape = concrete_shape(tensor);
    if tensor_shape == target_shape {
        return tensor.clone();
    }

    assert!(
        tensor.rank <= target_shape.len(),
        "cannot broadcast higher-rank shape {:?} to {:?}",
        tensor_shape,
        target_shape
    );

    let rank_diff = target_shape.len() - tensor.rank;
    let mut source_shape = Vec::with_capacity(tensor.rank);
    let mut dims = Vec::with_capacity(tensor.rank);
    for (src_axis, &src_dim) in tensor_shape.iter().enumerate() {
        let dst_axis = src_axis + rank_diff;
        let dst_dim = target_shape[dst_axis];
        assert!(
            src_dim == dst_dim || src_dim == 1,
            "cannot broadcast shape {:?} to {:?}",
            tensor_shape,
            target_shape
        );
        if src_dim == 1 && dst_dim != 1 {
            continue;
        }
        source_shape.push(src_dim);
        dims.push(dst_axis);
    }

    let source = if source_shape == tensor_shape {
        tensor.clone()
    } else {
        tensor.reshape(&source_shape)
    };
    source.broadcast_in_dim(target_shape, &dims)
}

/// Broadcast two tensors to a common shape.
fn broadcast_binary(a: &TracedTensor, b: &TracedTensor) -> (TracedTensor, TracedTensor) {
    if a.shape_hint == b.shape_hint && a.rank == b.rank {
        return (a.clone(), b.clone());
    }
    let a_shape = concrete_shape(a);
    let b_shape = concrete_shape(b);
    let target = broadcast_shape(&a_shape, &b_shape).unwrap_or_else(|| {
        panic!(
            "incompatible shapes for broadcast: {:?} and {:?}",
            a_shape, b_shape
        )
    });
    (broadcast_to(a, &target), broadcast_to(b, &target))
}

fn scale_with_constant(input: &TracedTensor, op: StdTensorOp) -> TracedTensor {
    let scalar = apply_nullary(op, 0, input.dtype, Some(vec![]));
    let input_shape = concrete_shape(input);
    let factor = broadcast_to(&scalar, &input_shape);
    apply_binary(
        StdTensorOp::Mul,
        input,
        &factor,
        input.rank,
        input.shape_hint.clone(),
    )
}

impl std::ops::Add for &TracedTensor {
    type Output = TracedTensor;

    fn add(self, rhs: &TracedTensor) -> TracedTensor {
        TracedTensor::add(self, rhs)
    }
}

impl std::ops::Mul for &TracedTensor {
    type Output = TracedTensor;

    fn mul(self, rhs: &TracedTensor) -> TracedTensor {
        TracedTensor::mul(self, rhs)
    }
}

impl std::ops::Mul<f64> for &TracedTensor {
    type Output = TracedTensor;

    fn mul(self, rhs: f64) -> TracedTensor {
        self.scale_real(rhs)
    }
}

impl std::ops::Mul<&TracedTensor> for f64 {
    type Output = TracedTensor;

    fn mul(self, rhs: &TracedTensor) -> TracedTensor {
        rhs.scale_real(self)
    }
}

impl std::ops::Neg for &TracedTensor {
    type Output = TracedTensor;

    fn neg(self) -> TracedTensor {
        TracedTensor::neg(self)
    }
}

impl std::ops::Div for &TracedTensor {
    type Output = TracedTensor;

    fn div(self, rhs: &TracedTensor) -> TracedTensor {
        TracedTensor::div(self, rhs)
    }
}

impl TracedTensor {
    /// Build a [`TracedTensor`] leaf from a concrete [`Tensor`], keeping its
    /// shape as a concrete `shape_hint`.
    ///
    /// This is the common constructor when you have concrete tensor data that
    /// you want to use both for graph building and for evaluation. The
    /// resulting tensor is treated as a concrete-shape leaf by downstream
    /// passes (binary einsum decomposition, build-time reshape folding, etc.).
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::{Tensor, TracedTensor};
    ///
    /// let a = TracedTensor::from_tensor_concrete_shape(
    ///     Tensor::from_vec(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]),
    /// );
    /// assert_eq!(a.rank, 2);
    /// assert!(a.is_concrete_shape());
    /// ```
    pub fn from_tensor_concrete_shape(tensor: Tensor) -> Self {
        let shape = tensor.shape().to_vec();
        let rank = shape.len();
        let dtype = tensor.dtype();
        let key = next_input_key();
        let data = Arc::new(tensor);

        let mut builder = FragmentBuilder::new();
        let val = builder.add_input(key.clone());
        builder.set_outputs(vec![val]);
        let fragment = Arc::new(builder.build());

        let mut map = HashMap::new();
        map.insert(key, Arc::clone(&data));

        Self {
            id: next_traced_id(),
            rank,
            dtype,
            fragment,
            val,
            data: Some(data),
            shape_hint: Some(shape.into_iter().map(SymDim::from).collect()),
            inputs_map: Arc::new(map),
            extra_roots: Vec::new(),
            checkpoint_chain: None,
        }
    }

    /// Build a [`TracedTensor`] leaf from a concrete [`Tensor`] but advertise
    /// a symbolic shape during graph construction.
    ///
    /// The tensor data is still attached (so plain `eval` works without
    /// bindings), but graph passes see the leaf as shape-symbolic. This is
    /// useful for building a single traced program that should not bake in
    /// shape-specific optimizations — e.g. mixing a known-shape tensor into
    /// an einsum with other `input_symbolic_shape` placeholders forces the
    /// einsum to be kept as a single `NaryEinsum` op rather than
    /// decomposing at build time.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::{Tensor, TracedTensor};
    ///
    /// let t = TracedTensor::from_tensor_symbolic_shape(
    ///     Tensor::from_vec(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]),
    /// );
    /// assert_eq!(t.rank, 2);
    /// assert!(!t.is_concrete_shape());
    /// ```
    pub fn from_tensor_symbolic_shape(tensor: Tensor) -> Self {
        let rank = tensor.shape().len();
        let dtype = tensor.dtype();
        let key = next_input_key();
        let data = Arc::new(tensor);

        let mut builder = FragmentBuilder::new();
        let val = builder.add_input(key.clone());
        builder.set_outputs(vec![val]);
        let fragment = Arc::new(builder.build());

        let mut map = HashMap::new();
        map.insert(key, Arc::clone(&data));

        Self {
            id: next_traced_id(),
            rank,
            dtype,
            fragment,
            val,
            data: Some(data),
            shape_hint: None,
            inputs_map: Arc::new(map),
            extra_roots: Vec::new(),
            checkpoint_chain: None,
        }
    }

    /// Build a data-less placeholder leaf with a fixed (concrete) shape.
    ///
    /// Must be bound via [`TracedTensor::eval_with_inputs`] before evaluation.
    /// Use this when you know the exact shape of the input but want to build
    /// the graph once and feed different concrete tensors at eval time.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_tensor::DType;
    /// use tenferro::TracedTensor;
    ///
    /// let x = TracedTensor::input_concrete_shape(DType::F64, &[2, 3]);
    /// assert_eq!(x.rank, 2);
    /// assert!(x.is_concrete_shape());
    /// ```
    pub fn input_concrete_shape(dtype: DType, shape: &[usize]) -> Self {
        let shape = shape.to_vec();
        let rank = shape.len();
        let key = next_input_key();

        let mut builder = FragmentBuilder::new();
        let val = builder.add_input(key.clone());
        builder.set_outputs(vec![val]);
        let fragment = Arc::new(builder.build());

        Self {
            id: next_traced_id(),
            rank,
            dtype,
            fragment,
            val,
            data: None,
            shape_hint: Some(shape.into_iter().map(SymDim::from).collect()),
            inputs_map: Arc::new(HashMap::new()),
            extra_roots: Vec::new(),
            checkpoint_chain: None,
        }
    }

    /// Build a data-less placeholder leaf with the given rank but fully
    /// symbolic shape (every dim is a distinct `SymDim::TensorAxis`).
    ///
    /// Must be bound via [`TracedTensor::eval_with_inputs`] before
    /// evaluation. Use this to build shape-agnostic graphs — in particular,
    /// einsum calls containing at least one `input_symbolic_shape` input are
    /// kept as a single `NaryEinsum` op so the contraction path can be
    /// optimized at eval time against the actual bound shapes.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_tensor::DType;
    /// use tenferro::TracedTensor;
    ///
    /// let x = TracedTensor::input_symbolic_shape(DType::F64, 2);
    /// assert_eq!(x.rank, 2);
    /// assert!(!x.is_concrete_shape());
    /// ```
    pub fn input_symbolic_shape(dtype: DType, rank: usize) -> Self {
        let key = next_input_key();

        let mut builder = FragmentBuilder::new();
        let val = builder.add_input(key.clone());
        builder.set_outputs(vec![val]);
        let fragment = Arc::new(builder.build());

        Self {
            id: next_traced_id(),
            rank,
            dtype,
            fragment,
            val,
            data: None,
            shape_hint: None,
            inputs_map: Arc::new(HashMap::new()),
            extra_roots: Vec::new(),
            checkpoint_chain: None,
        }
    }

    /// Build a concrete-shape [`TracedTensor`] leaf from typed `Vec<T>`
    /// data. Equivalent to
    /// [`TracedTensor::from_tensor_concrete_shape`]`(Tensor::from_vec(shape, data))`.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::TracedTensor;
    ///
    /// let a = TracedTensor::from_vec(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
    /// assert_eq!(a.rank, 2);
    /// ```
    pub fn from_vec<T: TensorScalar>(shape: Vec<usize>, data: Vec<T>) -> Self {
        Self::from_tensor_concrete_shape(T::into_tensor(shape, data))
    }

    /// Returns `true` iff every dim of this tensor's `shape_hint` is a
    /// constant `SymDim` (i.e. the shape is fully known at graph-build time).
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_tensor::DType;
    /// use tenferro::TracedTensor;
    ///
    /// let a = TracedTensor::from_vec(vec![2, 3], vec![1.0_f64; 6]);
    /// let b = TracedTensor::input_symbolic_shape(DType::F64, 2);
    /// assert!(a.is_concrete_shape());
    /// assert!(!b.is_concrete_shape());
    /// ```
    pub fn is_concrete_shape(&self) -> bool {
        try_concrete_shape(self).is_some()
    }

    /// If this `TracedTensor` is a leaf (single-node input fragment),
    /// return its input key. Computed tensors return `None`.
    pub fn input_key(&self) -> Option<TensorInputKey> {
        match &self.fragment.vals()[self.val].key {
            GlobalValKey::Input(key) => Some(key.clone()),
            _ => None,
        }
    }

    pub fn eval<B: TensorBackend>(&mut self, engine: &mut Engine<B>) -> Result<&Tensor> {
        self.eval_with_inputs(engine, &[])
    }

    /// Evaluate this traced tensor, binding external tensors to any
    /// placeholder leaves present in the graph.
    ///
    /// Each `(placeholder, tensor)` pair maps a placeholder
    /// [`TracedTensor`] (built via
    /// [`TracedTensor::input_concrete_shape`] or
    /// [`TracedTensor::input_symbolic_shape`]) to the concrete
    /// [`Tensor`] that should take its place during execution. Leaves that
    /// carry their own data (from [`TracedTensor::from_vec`],
    /// [`TracedTensor::from_tensor_concrete_shape`], or
    /// [`TracedTensor::from_tensor_symbolic_shape`]) must not appear in
    /// `bindings`.
    ///
    /// # Errors
    ///
    /// - [`Error::UnexpectedBinding`] if a binding's left side is not a
    ///   data-less placeholder leaf.
    /// - [`Error::DuplicateBinding`] if the same placeholder key appears
    ///   more than once in `bindings`.
    /// - [`Error::PlaceholderDtypeMismatch`] if a binding tensor's dtype
    ///   differs from the placeholder's declared dtype.
    /// - [`Error::PlaceholderShapeMismatch`] if a binding tensor's shape
    ///   differs from an `input_concrete_shape` placeholder's fixed shape.
    /// - [`Error::PlaceholderRankMismatch`] if a binding tensor's rank
    ///   differs from an `input_symbolic_shape` placeholder's rank.
    /// - [`Error::UnboundPlaceholder`] if the compiled graph contains a
    ///   placeholder that has no entry in `bindings`.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::{CpuBackend, Engine, Tensor, TracedTensor};
    /// use tenferro_tensor::DType;
    ///
    /// let mut engine = Engine::new(CpuBackend::new());
    /// let x = TracedTensor::input_symbolic_shape(DType::F64, 1);
    /// let mut y = &x + &x;
    /// let concrete = Tensor::from_vec(vec![3], vec![1.0_f64, 2.0, 3.0]);
    /// let out = y.eval_with_inputs(&mut engine, &[(&x, &concrete)]).unwrap();
    /// assert_eq!(out.shape(), &[3]);
    /// ```
    pub fn eval_with_inputs<B: TensorBackend>(
        &mut self,
        engine: &mut Engine<B>,
        bindings: &[(&TracedTensor, &Tensor)],
    ) -> Result<&Tensor> {
        // Build the binding map keyed on `TensorInputKey`, validating each
        // binding as we go. We validate *before* the already-evaluated
        // shortcut so that user mistakes (e.g. binding a data-carrying leaf)
        // surface as errors rather than getting silently ignored when the
        // output tensor happens to be cached.
        let mut binding_map: HashMap<TensorInputKey, &Tensor> = HashMap::new();
        for (index, (placeholder, tensor)) in bindings.iter().enumerate() {
            if placeholder.data.is_some() {
                return Err(Error::UnexpectedBinding {
                    binding_index: index,
                });
            }
            let key = placeholder.input_key().ok_or(Error::UnexpectedBinding {
                binding_index: index,
            })?;

            if placeholder.dtype != tensor.dtype() {
                return Err(Error::PlaceholderDtypeMismatch {
                    expected: placeholder.dtype,
                    actual: tensor.dtype(),
                });
            }

            match try_concrete_shape(placeholder) {
                Some(expected_shape) => {
                    if expected_shape.as_slice() != tensor.shape() {
                        return Err(Error::PlaceholderShapeMismatch {
                            expected: expected_shape,
                            actual: tensor.shape().to_vec(),
                        });
                    }
                }
                None => {
                    if placeholder.rank != tensor.shape().len() {
                        return Err(Error::PlaceholderRankMismatch {
                            expected: placeholder.rank,
                            actual: tensor.shape().len(),
                        });
                    }
                }
            }

            if binding_map.insert(key.clone(), *tensor).is_some() {
                return Err(Error::DuplicateBinding {
                    input_key: format!("{:?}", key),
                });
            }
        }

        // Already-evaluated shortcut — safe to take now that bindings passed
        // validation.
        if self.data.is_some() {
            return Ok(self.data.as_ref().unwrap().as_ref());
        }

        let output_key = self.fragment.vals()[self.val].key.clone();

        let view = resolve(self.resolve_roots());
        let graph = materialize_merge(&view, &[output_key]);
        let compiled = compile(&graph);

        let mut input_tensors = Vec::with_capacity(graph.inputs.len());
        let mut input_dtypes = Vec::with_capacity(graph.inputs.len());
        let mut input_shapes = Vec::with_capacity(graph.inputs.len());
        for key in &graph.inputs {
            match key {
                GlobalValKey::Input(k) => {
                    if let Some(tensor) = self.inputs_map.get(k) {
                        input_tensors.push(tensor.as_ref().clone());
                        input_dtypes.push(tensor.dtype());
                        input_shapes.push(DimExpr::from_concrete(tensor.shape()));
                    } else if let Some(bound) = binding_map.remove(k) {
                        input_tensors.push((*bound).clone());
                        input_dtypes.push(bound.dtype());
                        input_shapes.push(DimExpr::from_concrete(bound.shape()));
                    } else {
                        return Err(Error::UnboundPlaceholder {
                            input_key: format!("{:?}", k),
                        });
                    }
                }
                _ => {
                    return Err(Error::Internal(
                        "expected Input key in graph inputs".to_string(),
                    ));
                }
            }
        }
        let exec = compile_std_to_exec(&compiled, &input_dtypes, &input_shapes);

        let cached_exec = engine.get_or_compile(exec);
        let mut results = engine.eval_exec_ir(&cached_exec, input_tensors)?;
        if results.len() != 1 {
            return Err(Error::Internal(format!(
                "expected 1 output, got {}",
                results.len()
            )));
        }

        self.data = Some(Arc::new(results.remove(0)));
        Ok(self.data.as_ref().unwrap().as_ref())
    }

    pub fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor> {
        if self.rank != 0 {
            return Err(Error::NonScalarGrad {
                shape: error_shape_hint(self),
            });
        }

        let ones = ones_tensor(self.dtype, vec![]);
        let seed = TracedTensor::from_tensor_concrete_shape(ones);
        Ok(self.vjp(wrt, &seed))
    }

    /// Like [`grad`](Self::grad) but returns `None` when the scalar output does
    /// not depend on `wrt`.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let maybe_dx = loss.try_grad(&x)?;
    /// ```
    pub fn try_grad(&self, wrt: &TracedTensor) -> Result<Option<TracedTensor>> {
        if self.rank != 0 {
            return Err(Error::NonScalarGrad {
                shape: error_shape_hint(self),
            });
        }

        let ones = ones_tensor(self.dtype, vec![]);
        let seed = TracedTensor::from_tensor_concrete_shape(ones);
        Ok(self.try_vjp(wrt, &seed))
    }

    /// Evaluate this tensor and replace its graph with a concrete leaf.
    ///
    /// This keeps downstream forward evaluation rooted at the concrete value
    /// while retaining the original fragment chain for later reverse-mode AD.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::{CpuBackend, Engine, TracedTensor};
    ///
    /// let mut engine = Engine::new(CpuBackend::new());
    /// let x = TracedTensor::from_vec(vec![], vec![3.0_f64]);
    /// let mut y = &x * &x;
    /// y.checkpoint(&mut engine).unwrap();
    /// assert_eq!(y.eval(&mut engine).unwrap().shape(), &[] as &[usize]);
    /// ```
    pub fn checkpoint<B: TensorBackend>(&mut self, engine: &mut Engine<B>) -> Result<()> {
        self.eval(engine)?;
        let data = self
            .data
            .clone()
            .ok_or_else(|| Error::Internal("checkpoint eval did not populate data".to_string()))?;
        let concrete_shape_hint = Some(data.shape().iter().copied().map(SymDim::from).collect());

        let old_fragment = self.fragment.clone();
        let old_output_key = old_fragment.vals()[self.val].key.clone();
        let old_inputs = (*self.inputs_map).clone();

        let new_key = next_input_key();
        let mut builder = FragmentBuilder::new();
        let leaf_val = builder.add_input(new_key.clone());
        builder.set_outputs(vec![leaf_val]);
        let new_fragment = Arc::new(builder.build());

        let node = CheckpointNode {
            fragment: old_fragment,
            alias_key: new_key.clone(),
            alias_target: old_output_key,
            old_inputs,
            prev: self.checkpoint_chain.take(),
        };

        self.fragment = new_fragment;
        self.val = leaf_val;
        self.extra_roots.clear();
        self.shape_hint = concrete_shape_hint;
        self.checkpoint_chain = Some(Arc::new(node));

        let mut merged = HashMap::new();
        if let Some(chain) = &self.checkpoint_chain {
            merged.extend(chain.collect_inputs());
        }
        merged.insert(new_key, data);
        self.inputs_map = Arc::new(merged);

        Ok(())
    }

    pub fn jvp(&self, wrt: &TracedTensor, tangent: &TracedTensor) -> TracedTensor {
        self.try_jvp(wrt, tangent)
            .unwrap_or_else(|| panic!("jvp output is inactive for {:?}", leaf_input_key(wrt)))
    }

    /// Like [`jvp`](Self::jvp) but returns `None` when the output does not
    /// depend on `wrt` (i.e. the tangent is structurally zero).
    pub fn try_jvp(&self, wrt: &TracedTensor, tangent: &TracedTensor) -> Option<TracedTensor> {
        let wrt_input_key = leaf_input_key(wrt);
        let output_key = self.fragment.vals()[self.val].key.clone();
        let aliases = self
            .checkpoint_chain
            .as_ref()
            .map(|chain| chain.collect_aliases())
            .unwrap_or_default();
        let checkpoint_fragments = self
            .checkpoint_chain
            .as_ref()
            .map(|chain| chain.collect_fragments())
            .unwrap_or_default();
        let mut roots = self.resolve_roots();
        roots.extend(checkpoint_fragments.iter().cloned());
        let view = resolve(roots);
        let mut ad_ctx = ShapeGuardContext::default();
        let linear = differentiate(
            &view,
            std::slice::from_ref(&output_key),
            std::slice::from_ref(&wrt_input_key),
            next_pass_id(),
            &mut ad_ctx,
            &aliases,
        );
        let tangent_output = linear.tangent_outputs[0]?;
        let tangent_input_key = linear_input_key(&linear.fragment, linear.tangent_inputs[0].1);

        let mut inputs_map = (*self.inputs_map).clone();
        if let Some(chain) = &self.checkpoint_chain {
            inputs_map.extend(chain.collect_inputs());
        }
        inputs_map.insert(
            tangent_input_key,
            tangent
                .data
                .clone()
                .unwrap_or_else(|| panic!("jvp tangent must have concrete tensor data")),
        );

        let mut extra_roots = vec![self.fragment.clone()];
        extra_roots.extend(checkpoint_fragments);
        extra_roots.extend(self.extra_roots.iter().cloned());

        Some(TracedTensor {
            id: next_traced_id(),
            rank: self.rank,
            dtype: self.dtype,
            fragment: Arc::new(linear.fragment),
            val: tangent_output,
            data: None,
            shape_hint: self.shape_hint.clone(),
            inputs_map: Arc::new(inputs_map),
            extra_roots,
            checkpoint_chain: self.checkpoint_chain.clone(),
        })
    }

    pub fn vjp(&self, wrt: &TracedTensor, cotangent: &TracedTensor) -> TracedTensor {
        self.try_vjp(wrt, cotangent)
            .unwrap_or_else(|| panic!("vjp output is inactive for {:?}", leaf_input_key(wrt)))
    }

    fn try_vjp(&self, wrt: &TracedTensor, cotangent: &TracedTensor) -> Option<TracedTensor> {
        let wrt_input_key = leaf_input_key(wrt);
        let output_key = self.fragment.vals()[self.val].key.clone();
        let aliases = self
            .checkpoint_chain
            .as_ref()
            .map(|chain| chain.collect_aliases())
            .unwrap_or_default();
        let checkpoint_fragments = self
            .checkpoint_chain
            .as_ref()
            .map(|chain| chain.collect_fragments())
            .unwrap_or_default();
        let mut roots = self.resolve_roots();
        roots.extend(checkpoint_fragments.iter().cloned());
        let view = resolve(roots);
        let mut ad_ctx = ShapeGuardContext::default();
        let linear = differentiate(
            &view,
            std::slice::from_ref(&output_key),
            std::slice::from_ref(&wrt_input_key),
            next_pass_id(),
            &mut ad_ctx,
            &aliases,
        );
        let linear_tangent_input_ids: Vec<LocalValId> = linear
            .tangent_inputs
            .iter()
            .map(|(_, local_id)| *local_id)
            .collect();
        let transposed = transpose(&linear, &mut ad_ctx);
        let linear_fragment = Arc::new(linear.fragment);
        let cotangent_output = transposed.tangent_outputs[0]?;
        let cotangent_input_key =
            linear_input_key(&transposed.fragment, transposed.tangent_inputs[0].1);

        let mut inputs_map = (*self.inputs_map).clone();
        if let Some(chain) = &self.checkpoint_chain {
            inputs_map.extend(chain.collect_inputs());
        }
        inputs_map.insert(
            cotangent_input_key.clone(),
            cotangent
                .data
                .clone()
                .unwrap_or_else(|| panic!("vjp cotangent must have concrete tensor data")),
        );
        let zero_tangent = Arc::new(zeros_tensor(
            wrt.dtype,
            try_concrete_shape(wrt).unwrap_or_else(|| vec![0; wrt.rank]),
        ));
        for (_, local_id) in &transposed.tangent_inputs {
            let tangent_input_key = linear_input_key(&transposed.fragment, *local_id);
            if tangent_input_key != cotangent_input_key {
                inputs_map.insert(tangent_input_key, Arc::clone(&zero_tangent));
            }
        }
        for local_id in linear_tangent_input_ids {
            let tangent_input_key = linear_input_key(&linear_fragment, local_id);
            inputs_map.insert(tangent_input_key, Arc::clone(&zero_tangent));
        }

        let mut extra_roots = vec![self.fragment.clone(), linear_fragment];
        extra_roots.extend(checkpoint_fragments);
        extra_roots.extend(self.extra_roots.iter().cloned());

        Some(TracedTensor {
            id: next_traced_id(),
            rank: wrt.rank,
            dtype: wrt.dtype,
            fragment: Arc::new(transposed.fragment),
            val: cotangent_output,
            data: None,
            shape_hint: wrt.shape_hint.clone(),
            inputs_map: Arc::new(inputs_map),
            extra_roots,
            checkpoint_chain: self.checkpoint_chain.clone(),
        })
    }

    /// Elementwise addition with NumPy-style broadcasting.
    ///
    /// Prefer using the `+` operator when it reads naturally.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.add(&z);
    /// let y2 = &x + &z;
    /// ```
    pub fn add(&self, other: &TracedTensor) -> TracedTensor {
        let (lhs, rhs) = broadcast_binary(self, other);
        apply_binary(
            StdTensorOp::Add,
            &lhs,
            &rhs,
            lhs.rank,
            lhs.shape_hint.clone(),
        )
    }

    /// Elementwise multiplication with NumPy-style broadcasting.
    ///
    /// Prefer using the `*` operator when it reads naturally.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.mul(&z);
    /// let y2 = &x * &z;
    /// ```
    pub fn mul(&self, other: &TracedTensor) -> TracedTensor {
        let (lhs, rhs) = broadcast_binary(self, other);
        apply_binary(
            StdTensorOp::Mul,
            &lhs,
            &rhs,
            lhs.rank,
            lhs.shape_hint.clone(),
        )
    }

    /// Elementwise division with NumPy-style broadcasting.
    ///
    /// Prefer using the `/` operator when it reads naturally.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.div(&z);
    /// let y2 = &x / &z;
    /// ```
    pub fn div(&self, other: &TracedTensor) -> TracedTensor {
        let (lhs, rhs) = broadcast_binary(self, other);
        apply_binary(
            StdTensorOp::Div,
            &lhs,
            &rhs,
            lhs.rank,
            lhs.shape_hint.clone(),
        )
    }

    /// Elementwise negation.
    ///
    /// Prefer using the unary `-` operator when it reads naturally.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.neg();
    /// let y2 = -&x;
    /// ```
    pub fn neg(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Neg, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise complex conjugate.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.conj();
    /// ```
    pub fn conj(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Conj, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise absolute value.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.abs();
    /// ```
    pub fn abs(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Abs, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise sign.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.sign();
    /// ```
    pub fn sign(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Sign, self, self.rank, self.shape_hint.clone())
    }

    /// Scale by a real scalar: `y = factor * x`.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.scale_real(2.0);
    /// ```
    pub fn scale_real(&self, factor: f64) -> TracedTensor {
        let op = match self.dtype {
            DType::F64 => StdTensorOp::constant_f64(factor),
            DType::F32 => StdTensorOp::constant_f32(factor as f32),
            DType::C64 => StdTensorOp::constant_c64(Complex64::new(factor, 0.0)),
            DType::C32 => StdTensorOp::constant_c32(Complex32::new(factor as f32, 0.0)),
        };
        scale_with_constant(self, op)
    }

    /// Scale by a complex scalar: `y = factor * x`.
    ///
    /// This currently supports complex tensors only. For real scaling, prefer
    /// [`scale_real`](Self::scale_real).
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// use num_complex::Complex64;
    /// let y = x.scale_complex(Complex64::new(0.0, 1.0)); // multiply by i
    /// ```
    pub fn scale_complex(&self, factor: Complex64) -> TracedTensor {
        match self.dtype {
            DType::C64 => scale_with_constant(self, StdTensorOp::constant_c64(factor)),
            DType::C32 => scale_with_constant(
                self,
                StdTensorOp::constant_c32(Complex32::new(factor.re as f32, factor.im as f32)),
            ),
            DType::F32 | DType::F64 => {
                panic!(
                    "scale_complex only supports complex tensors; use scale_real for real tensors"
                )
            }
        }
    }

    /// Elementwise exponential.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.exp();
    /// ```
    pub fn exp(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Exp, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise natural logarithm.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.log();
    /// ```
    pub fn log(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Log, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise sine.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.sin();
    /// ```
    pub fn sin(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Sin, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise cosine.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.cos();
    /// ```
    pub fn cos(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Cos, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise hyperbolic tangent.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.tanh();
    /// ```
    pub fn tanh(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Tanh, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise square root.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.sqrt();
    /// ```
    pub fn sqrt(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Sqrt, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise reciprocal square root.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.rsqrt();
    /// ```
    pub fn rsqrt(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Rsqrt, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise power with NumPy-style broadcasting.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = base.pow(&exp);
    /// ```
    pub fn pow(&self, other: &TracedTensor) -> TracedTensor {
        let (lhs, rhs) = broadcast_binary(self, other);
        apply_binary(
            StdTensorOp::Pow,
            &lhs,
            &rhs,
            lhs.rank,
            lhs.shape_hint.clone(),
        )
    }

    /// Elementwise `exp(x) - 1`.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.expm1();
    /// ```
    pub fn expm1(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Expm1, self, self.rank, self.shape_hint.clone())
    }

    /// Elementwise `log(1 + x)`.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.log1p();
    /// ```
    pub fn log1p(&self) -> TracedTensor {
        apply_unary(StdTensorOp::Log1p, self, self.rank, self.shape_hint.clone())
    }

    /// Convert the tensor to a different dtype.
    ///
    /// For real-to-complex conversions this embeds the real values as
    /// `x + 0i`. For complex-to-real conversions this extracts the real part.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// use tenferro::DType;
    ///
    /// let y = x.convert(DType::C64);
    /// ```
    pub fn convert(&self, to: DType) -> TracedTensor {
        if self.dtype == to {
            return self.clone();
        }

        apply_unary_with_dtype(
            StdTensorOp::Convert {
                from: self.dtype,
                to,
            },
            self,
            self.rank,
            self.shape_hint.clone(),
            to,
        )
    }

    /// Generalized tensor contraction.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = a.dot_general(&b, config);
    /// ```
    pub fn dot_general(&self, other: &TracedTensor, config: DotGeneralConfig) -> TracedTensor {
        config
            .validate_ranks(self.rank, other.rank)
            .expect("DotGeneral config rank validation failed");
        config
            .validate_dims()
            .expect("DotGeneral config dimension validation failed");
        let lhs_free: Vec<usize> = (0..config.lhs_rank)
            .filter(|d| {
                !config.lhs_contracting_dims.contains(d) && !config.lhs_batch_dims.contains(d)
            })
            .collect();
        let rhs_free: Vec<usize> = (0..config.rhs_rank)
            .filter(|d| {
                !config.rhs_contracting_dims.contains(d) && !config.rhs_batch_dims.contains(d)
            })
            .collect();
        let out_rank = config.lhs_batch_dims.len() + lhs_free.len() + rhs_free.len();
        let out_shape_hint = match (&self.shape_hint, &other.shape_hint) {
            (Some(lhs_shape), Some(rhs_shape)) => {
                let mut out_shape = Vec::with_capacity(out_rank);
                for &d in &lhs_free {
                    out_shape.push(lhs_shape[d].clone());
                }
                for &d in &rhs_free {
                    out_shape.push(rhs_shape[d].clone());
                }
                for &d in &config.lhs_batch_dims {
                    out_shape.push(lhs_shape[d].clone());
                }
                Some(out_shape)
            }
            _ => None,
        };

        apply_binary(
            StdTensorOp::DotGeneral(config),
            self,
            other,
            out_rank,
            out_shape_hint,
        )
    }

    /// Sum over the given axes.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.reduce_sum(&[0]);
    /// let y2 = x.sum(&[0]);
    /// ```
    pub fn reduce_sum(&self, axes: &[usize]) -> TracedTensor {
        let out_shape_hint = self.shape_hint.as_ref().map(|shape| {
            (0..shape.len())
                .filter(|d| !axes.contains(d))
                .map(|d| shape[d].clone())
                .collect()
        });
        apply_unary(
            StdTensorOp::ReduceSum {
                axes: axes.to_vec(),
                input_shape: DimExpr::input_shape(0, self.rank),
            },
            self,
            self.rank - axes.len(),
            out_shape_hint,
        )
    }

    /// Reshape without changing element order.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.reshape(&[2, 2]);
    /// ```
    pub fn reshape(&self, shape: &[usize]) -> TracedTensor {
        apply_unary(
            StdTensorOp::Reshape {
                from_shape: DimExpr::input_shape(0, self.rank),
                to_shape: DimExpr::from_concrete(shape),
            },
            self,
            shape.len(),
            Some(shape.iter().copied().map(SymDim::from).collect()),
        )
    }

    /// Return a symbolic expression for the size of one axis.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let rows = x.sym_size(0);
    /// let cols = x.sym_size(1);
    /// let y = x.reshape_sym(&[rows * cols])?;
    /// ```
    pub fn sym_size(&self, axis: usize) -> SymDim {
        assert!(
            axis < self.rank,
            "axis {axis} out of bounds for rank {}",
            self.rank
        );
        self.shape_hint
            .as_ref()
            .and_then(|shape| shape.get(axis))
            .filter(|dim| dim.constant_value().is_none())
            .cloned()
            .unwrap_or_else(|| SymDim::tensor_axis(self.id, axis))
    }

    /// Reshape using symbolic dimensions derived from traced tensor axes.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let rows = x.sym_size(0);
    /// let cols = x.sym_size(1);
    /// let y = x.reshape_sym(&[rows * cols])?;
    /// ```
    pub fn reshape_sym(&self, shape: &[SymDim]) -> Result<TracedTensor> {
        let tensor_map = [(self.id, 0usize)];
        let to_shape = shape
            .iter()
            .map(|dim| dim.to_dim_expr(&tensor_map).map_err(Error::Internal))
            .collect::<Result<Vec<_>>>()?;
        let out_shape_hint = Some(shape.to_vec());
        Ok(apply_unary(
            StdTensorOp::Reshape {
                from_shape: DimExpr::input_shape(0, self.rank),
                to_shape,
            },
            self,
            shape.len(),
            out_shape_hint,
        ))
    }

    /// Broadcast into a larger shape with explicit dimension placement.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.broadcast_in_dim(&[2, 3], &[1]);
    /// let y2 = x.broadcast(&[2, 3], &[1]);
    /// ```
    pub fn broadcast_in_dim(&self, shape: &[usize], dims: &[usize]) -> TracedTensor {
        apply_unary(
            StdTensorOp::BroadcastInDim {
                shape: DimExpr::from_concrete(shape),
                dims: dims.to_vec(),
            },
            self,
            shape.len(),
            Some(shape.iter().copied().map(SymDim::from).collect()),
        )
    }

    /// Permute tensor axes.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.transpose(&[1, 0]);
    /// ```
    pub fn transpose(&self, perm: &[usize]) -> TracedTensor {
        let out_shape_hint = self
            .shape_hint
            .as_ref()
            .map(|shape| perm.iter().map(|&p| shape[p].clone()).collect());
        apply_unary(
            StdTensorOp::Transpose {
                perm: perm.to_vec(),
            },
            self,
            self.rank,
            out_shape_hint,
        )
    }

    /// Extract the diagonal along two axes.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.extract_diag(0, 1);
    /// ```
    pub fn extract_diag(&self, axis_a: usize, axis_b: usize) -> TracedTensor {
        assert!(
            axis_a < self.rank && axis_b < self.rank && axis_a != axis_b,
            "extract_diag: invalid axes"
        );
        let out_shape_hint = self.shape_hint.as_ref().map(|shape| {
            shape
                .iter()
                .enumerate()
                .filter_map(|(axis, dim)| (axis != axis_b).then_some(dim.clone()))
                .collect()
        });
        apply_unary(
            StdTensorOp::ExtractDiag { axis_a, axis_b },
            self,
            self.rank - 1,
            out_shape_hint,
        )
    }

    /// Embed a vector or lower-rank tensor along a diagonal.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.embed_diag(0, 1);
    /// ```
    pub fn embed_diag(&self, axis_a: usize, axis_b: usize) -> TracedTensor {
        assert!(
            axis_a < self.rank && axis_b <= self.rank,
            "embed_diag: invalid axes"
        );
        let out_shape_hint = self.shape_hint.as_ref().map(|shape| {
            let mut out_shape = shape.clone();
            out_shape.insert(axis_b, shape[axis_a].clone());
            out_shape
        });
        apply_unary(
            StdTensorOp::EmbedDiag { axis_a, axis_b },
            self,
            self.rank + 1,
            out_shape_hint,
        )
    }

    /// Alias for [`Self::reduce_sum`].
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.sum(&[0]);
    /// ```
    pub fn sum(&self, axes: &[usize]) -> TracedTensor {
        self.reduce_sum(axes)
    }

    /// Alias for [`Self::broadcast_in_dim`].
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let y = x.broadcast(&[2, 3], &[1]);
    /// ```
    pub fn broadcast(&self, shape: &[usize], dims: &[usize]) -> TracedTensor {
        self.broadcast_in_dim(shape, dims)
    }

    /// Return the runtime size of one axis as a scalar `f64` tensor.
    ///
    /// The result is metadata-derived and therefore has no gradient.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::{CpuBackend, Engine, TracedTensor};
    ///
    /// let mut engine = Engine::new(CpuBackend::new());
    /// let x = TracedTensor::from_vec(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
    /// let mut cols = x.shape_of(1);
    /// assert_eq!(cols.eval(&mut engine).unwrap().shape(), &[] as &[usize]);
    /// ```
    pub fn shape_of(&self, axis: usize) -> TracedTensor {
        assert!(
            axis < self.rank,
            "axis {axis} out of bounds for rank {}",
            self.rank
        );
        apply_unary_with_dtype(
            StdTensorOp::ShapeOf { axis },
            self,
            0,
            Some(vec![]),
            DType::F64,
        )
    }

    /// Truncate this tensor along `axis` to the first `size` elements.
    ///
    /// `size` is read at runtime from a scalar traced tensor. Values are
    /// rounded to the nearest integer, clamped to `[0, self.shape[axis]]`,
    /// and the output keeps the same element dtype as the input.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::{CpuBackend, Engine, TracedTensor};
    ///
    /// let mut engine = Engine::new(CpuBackend::new());
    /// let x = TracedTensor::from_vec(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]);
    /// let size = TracedTensor::from_vec(vec![], vec![2.0_f64]);
    /// let mut y = x.dynamic_truncate(&size, 0);
    /// assert_eq!(y.eval(&mut engine).unwrap().shape(), &[2]);
    /// ```
    pub fn dynamic_truncate(&self, size: &TracedTensor, axis: usize) -> TracedTensor {
        assert!(
            axis < self.rank,
            "axis {axis} out of bounds for rank {}",
            self.rank
        );
        assert!(
            size.rank == 0,
            "dynamic_truncate size must be a scalar tensor, got rank {}",
            size.rank
        );
        apply_binary(
            StdTensorOp::DynamicTruncate { axis },
            self,
            size,
            self.rank,
            None,
        )
    }

    /// Pad this tensor with zeros along `axis` to match `reference.shape[axis]`.
    ///
    /// If `reference` is smaller along that axis, this is a no-op.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro::{CpuBackend, Engine, TracedTensor};
    ///
    /// let mut engine = Engine::new(CpuBackend::new());
    /// let x = TracedTensor::from_vec(vec![2], vec![1.0_f64, 2.0]);
    /// let reference = TracedTensor::from_vec(vec![4], vec![0.0_f64, 0.0, 0.0, 0.0]);
    /// let mut y = x.pad_to_match(&reference, 0);
    /// assert_eq!(y.eval(&mut engine).unwrap().shape(), &[4]);
    /// ```
    pub fn pad_to_match(&self, reference: &TracedTensor, axis: usize) -> TracedTensor {
        assert!(
            axis < self.rank,
            "axis {axis} out of bounds for rank {}",
            self.rank
        );
        assert!(
            axis < reference.rank,
            "reference axis {axis} out of bounds for rank {}",
            reference.rank
        );
        apply_binary(
            StdTensorOp::PadToMatch { axis },
            self,
            reference,
            self.rank,
            reference.shape_hint.clone(),
        )
    }
}

pub(crate) fn apply_unary(
    op: StdTensorOp,
    input: &TracedTensor,
    out_rank: usize,
    out_shape_hint: Option<Vec<SymDim>>,
) -> TracedTensor {
    apply_unary_with_dtype(op, input, out_rank, out_shape_hint, input.dtype)
}

pub(crate) fn apply_unary_with_dtype(
    op: StdTensorOp,
    input: &TracedTensor,
    out_rank: usize,
    out_shape_hint: Option<Vec<SymDim>>,
    out_dtype: DType,
) -> TracedTensor {
    let mut builder = FragmentBuilder::new();
    builder.add_parent(input.fragment.clone());
    let input_ref = ValRef::External(input.fragment.vals()[input.val].key.clone());
    let outputs = builder.add_op(op, vec![input_ref], OpMode::Primal);
    builder.set_outputs(outputs.clone());
    let fragment = Arc::new(builder.build());

    TracedTensor {
        id: next_traced_id(),
        rank: out_rank,
        dtype: out_dtype,
        fragment,
        val: outputs[0],
        data: None,
        shape_hint: out_shape_hint,
        inputs_map: input.inputs_map.clone(),
        extra_roots: input.extra_roots.clone(),
        checkpoint_chain: input.checkpoint_chain.clone(),
    }
}

pub(crate) fn apply_nullary(
    op: StdTensorOp,
    rank: usize,
    dtype: DType,
    shape_hint: Option<Vec<SymDim>>,
) -> TracedTensor {
    let mut builder = FragmentBuilder::new();
    let outputs = builder.add_op(op, vec![], OpMode::Primal);
    builder.set_outputs(outputs.clone());
    let fragment = Arc::new(builder.build());

    TracedTensor {
        id: next_traced_id(),
        rank,
        dtype,
        fragment,
        val: outputs[0],
        data: None,
        shape_hint,
        inputs_map: Arc::new(HashMap::new()),
        extra_roots: Vec::new(),
        checkpoint_chain: None,
    }
}

pub(crate) fn apply_binary(
    op: StdTensorOp,
    lhs: &TracedTensor,
    rhs: &TracedTensor,
    out_rank: usize,
    out_shape_hint: Option<Vec<SymDim>>,
) -> TracedTensor {
    let mut builder = FragmentBuilder::new();
    builder.add_parent(lhs.fragment.clone());
    builder.add_parent(rhs.fragment.clone());
    let lhs_ref = ValRef::External(lhs.fragment.vals()[lhs.val].key.clone());
    let rhs_ref = ValRef::External(rhs.fragment.vals()[rhs.val].key.clone());
    let outputs = builder.add_op(op, vec![lhs_ref, rhs_ref], OpMode::Primal);
    builder.set_outputs(outputs.clone());
    let fragment = Arc::new(builder.build());

    let mut merged = (*lhs.inputs_map).clone();
    merged.extend(rhs.inputs_map.iter().map(|(k, v)| (k.clone(), v.clone())));
    let mut extra_roots = lhs.extra_roots.clone();
    extra_roots.extend(rhs.extra_roots.iter().cloned());

    TracedTensor {
        id: next_traced_id(),
        rank: out_rank,
        dtype: lhs.dtype,
        fragment,
        val: outputs[0],
        data: None,
        shape_hint: out_shape_hint,
        inputs_map: Arc::new(merged),
        extra_roots,
        checkpoint_chain: CheckpointNode::merge_chains(
            lhs.checkpoint_chain.clone(),
            rhs.checkpoint_chain.clone(),
        ),
    }
}

pub(crate) fn apply_multi_output(
    op: StdTensorOp,
    input: &TracedTensor,
    output_shapes: Vec<Vec<SymDim>>,
) -> Vec<TracedTensor> {
    let mut builder = FragmentBuilder::new();
    builder.add_parent(input.fragment.clone());
    let input_ref = ValRef::External(input.fragment.vals()[input.val].key.clone());
    let outputs = builder.add_op(op, vec![input_ref], OpMode::Primal);
    builder.set_outputs(outputs.clone());
    let fragment = Arc::new(builder.build());
    assert_eq!(
        outputs.len(),
        output_shapes.len(),
        "apply_multi_output: output count must match output_shapes"
    );

    outputs
        .iter()
        .zip(output_shapes)
        .map(|(&val, shape)| TracedTensor {
            id: next_traced_id(),
            rank: shape.len(),
            dtype: input.dtype,
            fragment: fragment.clone(),
            val,
            data: None,
            shape_hint: Some(shape),
            inputs_map: input.inputs_map.clone(),
            extra_roots: input.extra_roots.clone(),
            checkpoint_chain: input.checkpoint_chain.clone(),
        })
        .collect()
}

impl TracedTensor {
    fn resolve_roots(&self) -> Vec<Arc<Fragment<StdTensorOp>>> {
        let mut roots = Vec::with_capacity(1 + self.extra_roots.len());
        roots.push(self.fragment.clone());
        roots.extend(self.extra_roots.iter().cloned());
        roots
    }
}

fn leaf_input_key(tt: &TracedTensor) -> TensorInputKey {
    match &tt.fragment.vals()[tt.val].key {
        GlobalValKey::Input(key) => key.clone(),
        other => panic!("expected traced leaf input, got {:?}", other),
    }
}

fn linear_input_key(fragment: &Fragment<StdTensorOp>, local_id: LocalValId) -> TensorInputKey {
    match &fragment.vals()[local_id].key {
        GlobalValKey::Input(key) => key.clone(),
        other => panic!("expected linear fragment input, got {:?}", other),
    }
}

fn ones_tensor(dtype: DType, shape: Vec<usize>) -> Tensor {
    match dtype {
        DType::F32 => Tensor::F32(TypedTensor::ones(shape)),
        DType::F64 => Tensor::F64(TypedTensor::ones(shape)),
        DType::C32 => Tensor::C32(TypedTensor::ones(shape)),
        DType::C64 => Tensor::C64(TypedTensor::ones(shape)),
    }
}

fn zeros_tensor(dtype: DType, shape: Vec<usize>) -> Tensor {
    match dtype {
        DType::F32 => Tensor::F32(TypedTensor::zeros(shape)),
        DType::F64 => Tensor::F64(TypedTensor::zeros(shape)),
        DType::C32 => Tensor::C32(TypedTensor::zeros(shape)),
        DType::C64 => Tensor::C64(TypedTensor::zeros(shape)),
    }
}

pub fn eval_all<B: TensorBackend>(
    engine: &mut Engine<B>,
    outputs: &mut [&mut TracedTensor],
) -> Result<Vec<Tensor>> {
    let mut all_fragments = Vec::new();
    let mut output_keys = Vec::new();
    let mut all_inputs: HashMap<TensorInputKey, Arc<Tensor>> = HashMap::new();

    for tt in outputs.iter() {
        all_fragments.extend(tt.resolve_roots());
        output_keys.push(tt.fragment.vals()[tt.val].key.clone());
        all_inputs.extend(tt.inputs_map.iter().map(|(k, v)| (k.clone(), v.clone())));
    }

    let view = resolve(all_fragments);
    let graph = materialize_merge(&view, &output_keys);
    let compiled = compile(&graph);

    let mut input_tensors = Vec::with_capacity(graph.inputs.len());
    let mut input_dtypes = Vec::with_capacity(graph.inputs.len());
    let mut input_shapes = Vec::with_capacity(graph.inputs.len());
    for key in &graph.inputs {
        match key {
            GlobalValKey::Input(k) => {
                let tensor = all_inputs.get(k).ok_or_else(|| {
                    Error::MissingInput(format!("missing input data for key {:?}", k))
                })?;
                input_tensors.push(tensor.as_ref().clone());
                input_dtypes.push(tensor.dtype());
                input_shapes.push(DimExpr::from_concrete(tensor.shape()));
            }
            _ => {
                return Err(Error::Internal(
                    "expected Input key in graph inputs".to_string(),
                ));
            }
        }
    }
    let exec = compile_std_to_exec(&compiled, &input_dtypes, &input_shapes);

    let cached_exec = engine.get_or_compile(exec);
    let results: Vec<Tensor> = engine.eval_exec_ir(&cached_exec, input_tensors)?;

    for (tt, result) in outputs.iter_mut().zip(results.iter()) {
        tt.data = Some(Arc::new(result.clone()));
    }

    Ok(results)
}