tenferro-ad 0.8.0

use std::sync::Arc;

use num_complex::{Complex32, Complex64};
use tenferro_ops::dim_expr::DimExpr;
use tenferro_ops::ext_op::ExtensionOp;
use tenferro_ops::std_tensor_op::StdTensorOp;
use tenferro_runtime::program::{ProgramInputSpec, SemanticProgramBuilder};
use tenferro_runtime::{ErrorPhase, ExecutionEngine, GraphCompiler, Runtime};
use tenferro_tensor::DotGeneralConfig;
use tenferro_tensor::{DType, Engine, PadConfig, SliceConfig, Tensor, TensorRead, TypedTensor};

use crate::error::{Error, Result};
use crate::scalar_semantics::dynamic_truncate_size;
use crate::shape_infer::{promote_dtype, promote_dtype_for_binary_op, promote_dtypes};

enum PromotedTensorRead<'a> {
    Borrowed(Box<TensorRead<'a>>),
    Owned(Box<Tensor>),
}

enum ConcreteTensorRead<'a> {
    Borrowed(&'a Tensor),
    Owned(Box<Tensor>),
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum EagerInputPromotionPlan {
    Unchanged,
    All(DType),
    Pair {
        first: usize,
        second: usize,
        dtype: DType,
    },
}

impl EagerInputPromotionPlan {
    pub(crate) fn target_dtype(self, index: usize, original: DType) -> DType {
        match self {
            Self::Unchanged => original,
            Self::All(dtype) => dtype,
            Self::Pair {
                first,
                second,
                dtype,
            } if index == first || index == second => dtype,
            Self::Pair { .. } => original,
        }
    }
}

pub(crate) fn eager_input_promotion_plan(
    op: &StdTensorOp,
    input_count: usize,
    mut dtype_at: impl FnMut(usize) -> DType,
) -> EagerInputPromotionPlan {
    match op {
        StdTensorOp::Add
        | StdTensorOp::Sub
        | StdTensorOp::Mul
        | StdTensorOp::Div
        | StdTensorOp::Rem
        | StdTensorOp::Pow
        | StdTensorOp::Maximum
        | StdTensorOp::Minimum
        | StdTensorOp::Compare(_)
        | StdTensorOp::DotGeneral { .. } => EagerInputPromotionPlan::Pair {
            first: 0,
            second: 1,
            dtype: promote_dtype_for_binary_op(op, dtype_at(0), dtype_at(1)),
        },
        StdTensorOp::Select => EagerInputPromotionPlan::Pair {
            first: 1,
            second: 2,
            dtype: promote_dtype(dtype_at(1), dtype_at(2)),
        },
        StdTensorOp::Clamp | StdTensorOp::Concatenate { .. } => {
            EagerInputPromotionPlan::All(promote_dtypes((0..input_count).map(dtype_at)))
        }
        StdTensorOp::Scatter(_) => EagerInputPromotionPlan::Pair {
            first: 0,
            second: 2,
            dtype: promote_dtype(dtype_at(0), dtype_at(2)),
        },
        StdTensorOp::DynamicUpdateSlice => EagerInputPromotionPlan::Pair {
            first: 0,
            second: 1,
            dtype: promote_dtype(dtype_at(0), dtype_at(1)),
        },
        _ => EagerInputPromotionPlan::Unchanged,
    }
}

impl PromotedTensorRead<'_> {
    fn tensor_read(&self) -> TensorRead<'_> {
        match self {
            Self::Borrowed(read) => read.as_ref().clone(),
            Self::Owned(tensor) => TensorRead::from_tensor(tensor.as_ref()),
        }
    }
}

impl<'a> ConcreteTensorRead<'a> {
    fn tensor(&'a self) -> &'a Tensor {
        match self {
            Self::Borrowed(tensor) => tensor,
            Self::Owned(tensor) => tensor,
        }
    }
}

fn materialize_tensor_read(exec: &mut dyn Engine, input: TensorRead<'_>) -> Result<Tensor> {
    exec.to_contiguous_read(input).map_err(Error::from)
}

fn concrete_tensor_read<'a>(
    exec: &mut dyn Engine,
    input: TensorRead<'a>,
) -> Result<ConcreteTensorRead<'a>> {
    match input {
        TensorRead::Tensor(tensor) => Ok(ConcreteTensorRead::Borrowed(tensor)),
        TensorRead::View(view) => Ok(ConcreteTensorRead::Owned(Box::new(
            exec.to_contiguous_read(TensorRead::from_view(view))
                .map_err(Error::from)?,
        ))),
    }
}

fn concrete_tensor_reads<'a>(
    exec: &mut dyn Engine,
    inputs: &[TensorRead<'a>],
) -> Result<Vec<ConcreteTensorRead<'a>>> {
    inputs
        .iter()
        .cloned()
        .map(|input| concrete_tensor_read(exec, input))
        .collect()
}

fn promote_read_to_dtype<'a>(
    exec: &mut dyn Engine,
    input: TensorRead<'a>,
    promoted: DType,
) -> Result<PromotedTensorRead<'a>> {
    if input.dtype() == promoted {
        Ok(PromotedTensorRead::Borrowed(Box::new(input)))
    } else {
        let input = concrete_tensor_read(exec, input)?;
        Ok(PromotedTensorRead::Owned(Box::new(
            exec.convert(input.tensor(), promoted)
                .map_err(Error::from)?,
        )))
    }
}

fn promote_binary_reads_to_dtype<'a>(
    exec: &mut dyn Engine,
    a: TensorRead<'a>,
    b: TensorRead<'a>,
    promoted: DType,
) -> Result<(PromotedTensorRead<'a>, PromotedTensorRead<'a>)> {
    let a = promote_read_to_dtype(exec, a, promoted)?;
    let b = promote_read_to_dtype(exec, b, promoted)?;
    Ok((a, b))
}

fn promote_binary_reads<'a>(
    exec: &mut dyn Engine,
    a: TensorRead<'a>,
    b: TensorRead<'a>,
    op: &StdTensorOp,
) -> Result<(PromotedTensorRead<'a>, PromotedTensorRead<'a>)> {
    let plan = eager_input_promotion_plan(op, 2, |index| [a.dtype(), b.dtype()][index]);
    let promoted = plan.target_dtype(0, a.dtype());
    promote_binary_reads_to_dtype(exec, a, b, promoted)
}

pub(crate) fn exec_dot_general_with_conj_on_tensor_reads_in_session(
    lhs: TensorRead<'_>,
    rhs: TensorRead<'_>,
    config: &DotGeneralConfig,
    lhs_conj: bool,
    rhs_conj: bool,
    exec: &mut dyn Engine,
) -> Result<Tensor> {
    let op = StdTensorOp::DotGeneral {
        config: config.clone(),
    };
    let plan = eager_input_promotion_plan(&op, 2, |index| [lhs.dtype(), rhs.dtype()][index]);
    let promoted = plan.target_dtype(0, lhs.dtype());
    let (lhs, rhs) = promote_binary_reads_to_dtype(exec, lhs, rhs, promoted)?;
    exec.dot_general_with_conj_read(
        lhs.tensor_read(),
        rhs.tensor_read(),
        config,
        lhs_conj,
        rhs_conj,
    )
    .map_err(Error::from)
}

/// Execute a single [`StdTensorOp`] on concrete tensors.
///
/// This core helper rejects extension ops because they require a runtime owner
/// with an installed extension module.
#[cfg(test)]
pub(crate) fn exec_op_on_tensors(
    op: &StdTensorOp,
    inputs: &[&Tensor],
    engine: &mut dyn Engine,
) -> Result<Vec<Tensor>> {
    if let StdTensorOp::Extension(ext) = op {
        return Err(missing_extension_module_error(ext.as_ref()));
    }
    exec_standard_op_on_tensors_with_session(op, inputs, engine)
}

/// Run an extension op that has no prepared session executor: canonicalize the
/// reads on `engine`, then run the extension as a one-op compiled program on
/// the same engine. The engine kind must already be registered with `runtime`.
pub(crate) fn exec_extension_op_on_tensor_reads(
    op: &Arc<dyn ExtensionOp>,
    inputs: &[TensorRead<'_>],
    engine: &mut dyn ExecutionEngine,
    runtime: &Runtime,
) -> Result<Vec<Tensor>> {
    let concrete_inputs = concrete_tensor_reads(&mut *engine, inputs)?;
    let input_refs: Vec<&Tensor> = concrete_inputs.iter().map(|input| input.tensor()).collect();
    execute_extension_op_via_runtime(Arc::clone(op), &input_refs, engine, runtime)
}

fn execute_extension_op_via_runtime(
    op: Arc<dyn ExtensionOp>,
    inputs: &[&Tensor],
    engine: &mut dyn ExecutionEngine,
    runtime: &Runtime,
) -> Result<Vec<Tensor>> {
    ensure_extension_module_family(op.as_ref(), runtime)?;
    let mut builder = SemanticProgramBuilder::new();
    let input_values = inputs
        .iter()
        .map(|tensor| {
            let shape = tensor.shape().iter().copied().map(DimExpr::from);
            builder
                .input(ProgramInputSpec::new(tensor.dtype(), shape))
                .map_err(|source| {
                    Error::runtime_state_source(
                        "extension",
                        tenferro_runtime::ErrorPhase::Execution,
                        source,
                    )
                })
        })
        .collect::<Result<Vec<_>>>()?;
    let outputs = builder.add_extension(op, &input_values).map_err(|source| {
        Error::runtime_state_source("extension", tenferro_runtime::ErrorPhase::Execution, source)
    })?;
    let frozen = builder.finish(&outputs).map_err(|source| {
        Error::runtime_state_source("extension", tenferro_runtime::ErrorPhase::Execution, source)
    })?;
    let program = GraphCompiler::new().compile_frozen_program(&frozen)?;
    let reads = inputs
        .iter()
        .map(|tensor| TensorRead::from_tensor(tensor))
        .collect::<Vec<_>>();
    let target = engine.target();
    let prepared = runtime.prepare(&program, &reads, std::slice::from_ref(&target))?;
    crate::eager::run_prepared_on(runtime, engine, &prepared, &reads)
}

fn ensure_extension_module_family(op: &dyn ExtensionOp, runtime: &Runtime) -> Result<()> {
    let snapshot = runtime.snapshot().map_err(|source| {
        Error::runtime_state_source("extension", ErrorPhase::Execution, source)
    })?;
    if snapshot.has_extension_family(op.family_id()) {
        Ok(())
    } else {
        Err(missing_extension_module_error(op))
    }
}

fn missing_extension_module_error(ext: &dyn tenferro_ops::ext_op::ExtensionOp) -> Error {
    Error::TensorRuntime(tenferro_tensor::Error::invalid_argument(
        "extension",
        "extension_module",
        format!(
            "missing extension module for op family_id {:?}; install an ExtensionModule on the Runtime",
            ext.family_id()
        ),
    ))
}

fn axis_out_of_bounds(op: &'static str, axis: usize, rank: usize) -> Error {
    Error::TensorRuntime(tenferro_tensor::Error::axis_out_of_bounds(op, axis, rank))
}

fn invalid_config(op: &'static str, message: impl Into<String>) -> Error {
    Error::TensorRuntime(tenferro_tensor::Error::invalid_argument(
        op,
        "configuration",
        message,
    ))
}

fn pad_to_match_high_padding(target_size: usize, current_size: usize) -> Result<i64> {
    let delta = target_size
        .checked_sub(current_size)
        .ok_or_else(|| invalid_config("PadToMatch", "target size is smaller than current size"))?;
    i64::try_from(delta).map_err(|_| {
        invalid_config(
            "PadToMatch",
            format!("padding width {delta} does not fit in i64"),
        )
    })
}

fn shape_of_host_tensor(axis: usize, shape: &[usize]) -> Result<Tensor> {
    if axis >= shape.len() {
        return Err(axis_out_of_bounds("ShapeOf", axis, shape.len()));
    }
    let size = shape[axis] as f64;
    Ok(Tensor::from_typed::<f64>(TypedTensor::from_vec_col_major(
        vec![],
        vec![size],
    )?))
}

fn generated_host_output(
    op: &StdTensorOp,
    input_shape: Option<&[usize]>,
) -> Result<Option<Tensor>> {
    match op {
        StdTensorOp::Constant { dtype, bytes } => constant_tensor(*dtype, bytes).map(Some),
        StdTensorOp::ShapeOf { axis } => {
            let shape = input_shape.ok_or_else(|| invalid_config("ShapeOf", "missing input"))?;
            shape_of_host_tensor(*axis, shape).map(Some)
        }
        _ => Ok(None),
    }
}

pub(crate) fn exec_standard_op_on_tensor_reads_with_session(
    op: &StdTensorOp,
    inputs: &[TensorRead<'_>],
    exec: &mut dyn Engine,
) -> Result<Vec<Tensor>> {
    if let Some(host) = generated_host_output(op, inputs.first().map(TensorRead::shape))? {
        return exec
            .upload_host_tensor(TensorRead::from_tensor(&host))
            .map(|output| vec![output])
            .map_err(Error::from);
    }
    exec_standard_op_on_tensor_reads_in_session(op, inputs, exec)
}

pub(crate) fn exec_standard_op_on_tensors_with_session(
    op: &StdTensorOp,
    inputs: &[&Tensor],
    exec: &mut dyn Engine,
) -> Result<Vec<Tensor>> {
    if let Some(host) = generated_host_output(op, inputs.first().map(|tensor| tensor.shape()))? {
        return exec
            .upload_host_tensor(TensorRead::from_tensor(&host))
            .map(|output| vec![output])
            .map_err(Error::from);
    }
    exec_standard_op_on_tensors_in_session(op, inputs, exec)
}

pub(crate) fn exec_standard_op_on_tensor_reads_in_session(
    op: &StdTensorOp,
    inputs: &[TensorRead<'_>],
    exec: &mut dyn Engine,
) -> Result<Vec<Tensor>> {
    // `_read` hooks own layout handling, including backend-specific fallback.
    // Materialize only where the value itself is needed (`Convert`, the
    // dynamic-truncate size, and an operation that must return a fresh owned
    // tensor), never to prepare an operand for a hook (#2223).
    let promotion_plan =
        eager_input_promotion_plan(op, inputs.len(), |index| inputs[index].dtype());
    let target_dtype = |index| promotion_plan.target_dtype(index, inputs[index].dtype());
    let result = match op {
        StdTensorOp::Add => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.add_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Sub => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.sub_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Mul => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.mul_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Neg => vec![exec.neg_read(inputs[0].clone())?],
        StdTensorOp::Div => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.div_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Rem => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.rem_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Exp => vec![exec.exp_read(inputs[0].clone())?],
        StdTensorOp::Log => vec![exec.log_read(inputs[0].clone())?],
        StdTensorOp::Sin => vec![exec.sin_read(inputs[0].clone())?],
        StdTensorOp::Cos => vec![exec.cos_read(inputs[0].clone())?],
        StdTensorOp::Tanh => vec![exec.tanh_read(inputs[0].clone())?],
        StdTensorOp::Sqrt => vec![exec.sqrt_read(inputs[0].clone())?],
        StdTensorOp::Rsqrt => vec![exec.rsqrt_read(inputs[0].clone())?],
        StdTensorOp::Pow => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.pow_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Abs => vec![exec.abs_read(inputs[0].clone())?],
        StdTensorOp::Sign => vec![exec.sign_read(inputs[0].clone())?],
        StdTensorOp::Conj => vec![exec.conj_read(inputs[0].clone())?],
        StdTensorOp::Maximum => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.maximum_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Minimum => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.minimum_read(a.tensor_read(), b.tensor_read())?]
        }
        StdTensorOp::Compare(dir) => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            vec![exec.compare_read(a.tensor_read(), b.tensor_read(), dir)?]
        }
        StdTensorOp::Transpose { perm } => vec![exec.transpose_read(inputs[0].clone(), perm)?],
        StdTensorOp::ReduceSum { axes, .. } => {
            vec![exec.reduce_sum_read(inputs[0].clone(), axes)?]
        }
        StdTensorOp::ReduceSumSquares { axes, .. } => {
            vec![exec.reduce_sum_squares_read(inputs[0].clone(), axes)?]
        }
        StdTensorOp::DotGeneral { config, .. } => {
            let (a, b) = promote_binary_reads(exec, inputs[0].clone(), inputs[1].clone(), op)?;
            // One rule, shared with the compiler pass: a degenerate form is the
            // elementary op it actually is, not a contraction (invariant I4), so
            // the eager path does not pay the backend's contraction machinery
            // for it either.
            match tenferro_tensor::degenerate_dot_kind(
                config,
                a.tensor_read().shape(),
                b.tensor_read().shape(),
            ) {
                Some(tenferro_tensor::DegenerateDotKind::Hadamard) => {
                    vec![exec.mul_read(a.tensor_read(), b.tensor_read())?]
                }
                Some(tenferro_tensor::DegenerateDotKind::FullContraction) => {
                    let product = exec.mul_read(a.tensor_read(), b.tensor_read())?;
                    let axes: Vec<usize> = (0..product.shape().len()).collect();
                    vec![exec.reduce_sum_read(TensorRead::from_tensor(&product), &axes)?]
                }
                _ => vec![exec.dot_general_read(a.tensor_read(), b.tensor_read(), config)?],
            }
        }
        StdTensorOp::Reshape { to_shape, .. } => {
            let shape = resolve_tensor_read_shape_exprs(inputs, to_shape)?;
            vec![exec.reshape_read(inputs[0].clone(), &shape)?]
        }
        StdTensorOp::BroadcastInDim { shape, dims } => {
            let shape = resolve_tensor_read_shape_exprs(inputs, shape)?;
            vec![exec.broadcast_in_dim_read(inputs[0].clone(), &shape, dims)?]
        }
        // Structural and indexing operands stay reads: the backend reads the
        // operand's own layout when it can, and its `_read` default stages the
        // view when it cannot. Pre-materializing here would take that decision
        // away from the backend (#2223).
        StdTensorOp::ExtractDiag { axis_a, axis_b } => {
            vec![exec.extract_diagonal_read(inputs[0].clone(), *axis_a, *axis_b)?]
        }
        StdTensorOp::EmbedDiag { axis_a, axis_b } => {
            vec![exec.embed_diagonal_read(inputs[0].clone(), *axis_a, *axis_b)?]
        }
        StdTensorOp::Tril { k } => {
            vec![exec.tril_read(inputs[0].clone(), *k)?]
        }
        StdTensorOp::Triu { k } => {
            vec![exec.triu_read(inputs[0].clone(), *k)?]
        }
        StdTensorOp::Slice(config) => {
            vec![exec.slice_read(inputs[0].clone(), config)?]
        }
        StdTensorOp::Pad(config) => {
            vec![exec.pad_read(inputs[0].clone(), config)?]
        }
        StdTensorOp::Reverse { axes } => {
            vec![exec.reverse_read(inputs[0].clone(), axes)?]
        }
        StdTensorOp::ReduceProd { axes, .. } => {
            vec![exec.reduce_prod_read(inputs[0].clone(), axes)?]
        }
        StdTensorOp::ReduceMax { axes, .. } => {
            vec![exec.reduce_max_read(inputs[0].clone(), axes)?]
        }
        StdTensorOp::ReduceMin { axes, .. } => {
            vec![exec.reduce_min_read(inputs[0].clone(), axes)?]
        }
        StdTensorOp::Expm1 => vec![exec.expm1_read(inputs[0].clone())?],
        StdTensorOp::Log1p => vec![exec.log1p_read(inputs[0].clone())?],
        StdTensorOp::Erf => vec![exec.erf_read(inputs[0].clone())?],
        StdTensorOp::Convert { to, .. } => {
            let input = concrete_tensor_read(exec, inputs[0].clone())?;
            vec![exec.cast(input.tensor(), *to)?]
        }
        StdTensorOp::Constant { .. } => {
            return Err(Error::Internal(
                    "standard graph eager execution does not support Constant inside an existing backend session"
                        .into(),
                ));
        }
        StdTensorOp::Select => {
            let b = promote_read_to_dtype(exec, inputs[1].clone(), target_dtype(1))?;
            let c = promote_read_to_dtype(exec, inputs[2].clone(), target_dtype(2))?;
            vec![exec.select_read(inputs[0].clone(), b.tensor_read(), c.tensor_read())?]
        }
        StdTensorOp::Clamp => {
            let input = promote_read_to_dtype(exec, inputs[0].clone(), target_dtype(0))?;
            let lower = promote_read_to_dtype(exec, inputs[1].clone(), target_dtype(1))?;
            let upper = promote_read_to_dtype(exec, inputs[2].clone(), target_dtype(2))?;
            vec![exec.clamp_read(
                input.tensor_read(),
                lower.tensor_read(),
                upper.tensor_read(),
            )?]
        }
        StdTensorOp::Concatenate { axis, .. } => {
            let mut promoted = Vec::with_capacity(inputs.len());
            for (index, input) in inputs.iter().enumerate() {
                promoted.push(promote_read_to_dtype(
                    exec,
                    input.clone(),
                    target_dtype(index),
                )?);
            }
            let reads: Vec<TensorRead<'_>> = promoted
                .iter()
                .map(PromotedTensorRead::tensor_read)
                .collect();
            vec![exec.concatenate_read(&reads, *axis)?]
        }
        StdTensorOp::Gather(config) => {
            // Both operands stay reads: the backend decides whether it can
            // consume them in place (invariant I1).
            vec![exec.gather_read(inputs[0].clone(), inputs[1].clone(), config)?]
        }
        StdTensorOp::GatherDynamicSliceSizes {
            offset_dims,
            collapsed_slice_dims,
            start_index_map,
            index_vector_dim,
            slice_sizes,
        } => {
            let slice_sizes = resolve_tensor_read_shape_exprs(inputs, slice_sizes)?;
            let config = tenferro_tensor::GatherConfig {
                offset_dims: offset_dims.clone(),
                collapsed_slice_dims: collapsed_slice_dims.clone(),
                start_index_map: start_index_map.clone(),
                index_vector_dim: *index_vector_dim,
                slice_sizes,
            };
            vec![exec.gather_read(inputs[0].clone(), inputs[1].clone(), &config)?]
        }
        StdTensorOp::Scatter(config) => {
            let operand = promote_read_to_dtype(exec, inputs[0].clone(), target_dtype(0))?;
            let updates = promote_read_to_dtype(exec, inputs[2].clone(), target_dtype(2))?;
            vec![exec.scatter_read(
                operand.tensor_read(),
                inputs[1].clone(),
                updates.tensor_read(),
                config,
            )?]
        }
        StdTensorOp::DynamicSlice { slice_sizes } => {
            vec![exec.dynamic_slice_read(inputs[0].clone(), inputs[1].clone(), slice_sizes)?]
        }
        StdTensorOp::DynamicUpdateSlice => {
            let operand = promote_read_to_dtype(exec, inputs[0].clone(), target_dtype(0))?;
            let update = promote_read_to_dtype(exec, inputs[1].clone(), target_dtype(1))?;
            vec![exec.dynamic_update_slice_read(
                operand.tensor_read(),
                update.tensor_read(),
                inputs[2].clone(),
            )?]
        }
        StdTensorOp::ShapeOf { .. } => {
            return Err(Error::Internal(
                    "standard graph eager execution does not support ShapeOf inside an existing backend session"
                        .into(),
                ));
        }
        StdTensorOp::DynamicTruncate { axis } => {
            let input = &inputs[0];
            if *axis >= input.shape().len() {
                return Err(axis_out_of_bounds(
                    "DynamicTruncate",
                    *axis,
                    input.shape().len(),
                ));
            }
            let size_tensor = concrete_tensor_read(exec, inputs[1].clone())?;
            let axis_extent = input.shape()[*axis];
            let size = dynamic_truncate_size(size_tensor.tensor(), axis_extent)?;
            let rank = input.shape().len();
            let mut limits = input.shape().to_vec();
            limits[*axis] = size;
            let config = SliceConfig {
                starts: vec![0; rank],
                limits,
                strides: vec![1; rank],
            };
            vec![exec.slice_read(inputs[0].clone(), &config)?]
        }
        StdTensorOp::PadToMatch { axis } => {
            let input = &inputs[0];
            let reference = &inputs[1];
            if *axis >= input.shape().len() {
                return Err(axis_out_of_bounds("PadToMatch", *axis, input.shape().len()));
            }
            if *axis >= reference.shape().len() {
                return Err(axis_out_of_bounds(
                    "PadToMatch",
                    *axis,
                    reference.shape().len(),
                ));
            }
            let target_size = reference.shape()[*axis];
            let current_size = input.shape()[*axis];
            if current_size >= target_size {
                vec![materialize_tensor_read(exec, inputs[0].clone())?]
            } else {
                let rank = input.shape().len();
                let mut high = vec![0i64; rank];
                high[*axis] = pad_to_match_high_padding(target_size, current_size)?;
                let config = PadConfig {
                    edge_padding_low: vec![0i64; rank],
                    edge_padding_high: high,
                    interior_padding: vec![0i64; rank],
                };
                vec![exec.pad_read(inputs[0].clone(), &config)?]
            }
        }
        StdTensorOp::Extension(ext) => {
            return Err(missing_extension_module_error(ext.as_ref()));
        }
    };
    Ok(result)
}

pub(crate) fn exec_standard_op_on_tensors_in_session(
    op: &StdTensorOp,
    inputs: &[&Tensor],
    exec: &mut dyn Engine,
) -> Result<Vec<Tensor>> {
    // One dispatch table: owned inputs are reads of their tensors, so the owned
    // and borrowed entry points cannot drift apart.
    let reads: Vec<TensorRead<'_>> = inputs
        .iter()
        .map(|tensor| TensorRead::from_tensor(tensor))
        .collect();
    exec_standard_op_on_tensor_reads_in_session(op, &reads, exec)
}

fn resolve_tensor_read_shape_exprs(
    inputs: &[TensorRead<'_>],
    exprs: &[DimExpr],
) -> Result<Vec<usize>> {
    let input_shapes: Vec<&[usize]> = inputs.iter().map(|tensor| tensor.shape()).collect();
    DimExpr::eval_all(exprs, &input_shapes).map_err(|err| Error::ShapeExpressionEvaluation {
        expression: format!("{exprs:?}"),
        cause: err.into(),
    })
}

fn constant_tensor(dtype: DType, bytes: &[u8]) -> Result<Tensor> {
    Ok(match dtype {
        DType::F64 => Tensor::from_typed::<f64>(TypedTensor::from_vec_col_major(
            vec![],
            vec![f64::from_le_bytes(exact_bytes::<8>(dtype, bytes)?)],
        )?),
        DType::F32 => Tensor::from_typed::<f32>(TypedTensor::from_vec_col_major(
            vec![],
            vec![f32::from_le_bytes(exact_bytes::<4>(dtype, bytes)?)],
        )?),
        DType::I32 => Tensor::from_typed::<i32>(TypedTensor::from_vec_col_major(
            vec![],
            vec![i32::from_le_bytes(exact_bytes::<4>(dtype, bytes)?)],
        )?),
        DType::I64 => Tensor::from_typed::<i64>(TypedTensor::from_vec_col_major(
            vec![],
            vec![i64::from_le_bytes(exact_bytes::<8>(dtype, bytes)?)],
        )?),
        DType::Bool => Tensor::from_typed::<bool>(TypedTensor::from_vec_col_major(
            vec![],
            vec![exact_bytes::<1>(dtype, bytes)?[0] != 0],
        )?),
        DType::C64 => {
            let data = exact_bytes::<16>(dtype, bytes)?;
            let mut re_bytes = [0u8; 8];
            let mut im_bytes = [0u8; 8];
            re_bytes.copy_from_slice(&data[..8]);
            im_bytes.copy_from_slice(&data[8..]);
            let re = f64::from_le_bytes(re_bytes);
            let im = f64::from_le_bytes(im_bytes);
            Tensor::from_typed::<tenferro_tensor::Complex64>(TypedTensor::from_vec_col_major(
                vec![],
                vec![Complex64::new(re, im)],
            )?)
        }
        DType::C32 => {
            let data = exact_bytes::<8>(dtype, bytes)?;
            let mut re_bytes = [0u8; 4];
            let mut im_bytes = [0u8; 4];
            re_bytes.copy_from_slice(&data[..4]);
            im_bytes.copy_from_slice(&data[4..]);
            let re = f32::from_le_bytes(re_bytes);
            let im = f32::from_le_bytes(im_bytes);
            Tensor::from_typed::<tenferro_tensor::Complex32>(TypedTensor::from_vec_col_major(
                vec![],
                vec![Complex32::new(re, im)],
            )?)
        }
        // An externally defined scalar has no runtime constant representation, so
        // the eager constant path rejects it rather than decoding a guessed layout.
        DType::External(_) => {
            return Err(Error::TensorRuntime(
                tenferro_tensor::Error::invalid_argument(
                    "constant_tensor",
                    "dtype",
                    "an externally defined scalar has no runtime constant",
                ),
            ));
        }
    })
}

fn exact_bytes<const N: usize>(dtype: DType, bytes: &[u8]) -> Result<[u8; N]> {
    if bytes.len() != N {
        return Err(invalid_config(
            "Constant",
            format!(
                "constant {:?} expected {} bytes, got {}",
                dtype,
                N,
                bytes.len()
            ),
        ));
    }
    let mut out = [0u8; N];
    out.copy_from_slice(bytes);
    Ok(out)
}

#[cfg(test)]
mod tests;