tenferro-cpu 0.8.0

use std::mem::size_of_val;

use smallvec::SmallVec;
use strided_kernel::{
    col_major_strides, ErasedConcatenatePlan, ErasedDynamicSlicePlan, ErasedDynamicUpdateSlicePlan,
    ErasedGatherPlan, ErasedPadPlan, ErasedRawStridedPtr, ErasedRawStridedRef, ErasedReversePlan,
    ErasedScatterPlan, ErasedSlicePlan, Exec, GatherSpec, KernelDType, ScatterSpec,
};
use tinyvec::TinyVec;

use super::indexed_plan_cache::{IndexedPlanFamily, IndexedPlanKey};
use super::operand::{operand_from_view, Operand};
use super::typed_host_data;
use super::PooledUninitOutput;
use crate::buffer_pool::{BufferPool, PoolScalar};
use crate::CpuPlanCache;
use tenferro_tensor::TensorScalar;
use tenferro_tensor::{DType, GatherConfig, PadConfig, ScatterConfig, SliceConfig};
use tenferro_tensor::{Tensor, TensorRead, TensorView, TypedTensor, TypedTensorView};

type InlineStrides = TinyVec<[isize; 8]>;

fn inline_col_major_strides(op: &'static str, dims: &[usize]) -> crate::Result<InlineStrides> {
    let mut strides: InlineStrides = std::iter::repeat_n(1isize, dims.len()).collect();
    for axis in 1..dims.len() {
        if strides[axis - 1] == 0 {
            strides[axis] = 0;
            continue;
        }
        let prior_dim = isize::try_from(dims[axis - 1]).map_err(|_| {
            crate::Error::invalid_argument(
                op,
                "shape",
                format!(
                    "dimension {} does not fit in isize while computing column-major strides",
                    dims[axis - 1]
                ),
            )
        })?;
        strides[axis] = strides[axis - 1].checked_mul(prior_dim).ok_or_else(|| {
            crate::Error::invalid_argument(
                op,
                "shape",
                format!("column-major stride overflows isize for shape {dims:?}"),
            )
        })?;
    }
    Ok(strides)
}

// Indexing families delegate bulk traversal to strided-kernel erased plans.
// Engine entrypoints inject the engine-derived execution policy.

trait TensorAsTyped<T> {
    fn as_typed(&self) -> Option<&TypedTensor<T>>;
}

macro_rules! impl_tensor_as_typed {
    ($(($ty:ty, $variant:ident)),+ $(,)?) => {
        $(
            impl TensorAsTyped<$ty> for Tensor {
                fn as_typed(&self) -> Option<&TypedTensor<$ty>> {
                    let tensor = tenferro_tensor::Tensor::as_typed::<$ty>(self);
                    tensor
                }
            }
        )+
    };
}

impl_tensor_as_typed!(
    (f32, F32),
    (f64, F64),
    (i32, I32),
    (i64, I64),
    (bool, Bool),
    (num_complex::Complex<f32>, C32),
    (num_complex::Complex<f64>, C64),
);

fn kernel_dtype(dtype: DType) -> KernelDType {
    match dtype {
        DType::F32 => KernelDType::F32,
        DType::F64 => KernelDType::F64,
        DType::I32 => KernelDType::I32,
        DType::I64 => KernelDType::I64,
        DType::Bool => KernelDType::Bool,
        DType::C32 => KernelDType::C32,
        DType::C64 => KernelDType::C64,
        // INVARIANT: `KernelDType` is the fixed compiled-kernel vocabulary, and
        // callers reach this after rejecting unsupported dtypes. An externally
        // defined scalar has no entry in it.
        DType::External(_) => unreachable!("KernelDType covers the preset scalars"),
    }
}

fn typed_bytes<T>(data: &[T]) -> &[u8] {
    // SAFETY: `data` is an aligned typed slice. The returned byte slice has
    // the same lifetime and exact byte length, and is read-only.
    unsafe { std::slice::from_raw_parts(data.as_ptr().cast::<u8>(), size_of_val(data)) }
}

fn gather_spec(config: &GatherConfig) -> GatherSpec {
    GatherSpec {
        offset_dims: config.offset_dims.clone(),
        collapsed_slice_dims: config.collapsed_slice_dims.clone(),
        start_index_map: config.start_index_map.clone(),
        index_vector_dim: config.index_vector_dim,
        slice_sizes: config.slice_sizes.clone(),
    }
}

fn scatter_spec(config: &ScatterConfig) -> ScatterSpec {
    ScatterSpec {
        update_window_dims: config.update_window_dims.clone(),
        inserted_window_dims: config.inserted_window_dims.clone(),
        scatter_dims_to_operand_dims: config.scatter_dims_to_operand_dims.clone(),
        index_vector_dim: config.index_vector_dim,
    }
}

fn indexed_plan_key(
    family: IndexedPlanFamily,
    dtype: KernelDType,
    index_dtype: KernelDType,
    dims: &[&[usize]],
    strides: &[&[isize]],
    config: &[&[usize]],
) -> IndexedPlanKey {
    IndexedPlanKey::from_slices(family, dtype, index_dtype, dims, strides, config)
}

macro_rules! dispatch_tensor_view_result {
    ($view:expr, |$v:ident| $body:expr) => {
        match $view {
            TensorView::F32($v) => Ok(Tensor::from_typed::<f32>($body?)),
            TensorView::F64($v) => Ok(Tensor::from_typed::<f64>($body?)),
            TensorView::I32($v) => Ok(Tensor::from_typed::<i32>($body?)),
            TensorView::I64($v) => Ok(Tensor::from_typed::<i64>($body?)),
            TensorView::Bool($v) => Ok(Tensor::from_typed::<bool>($body?)),
            TensorView::C32($v) => Ok(Tensor::from_typed::<tenferro_tensor::Complex32>($body?)),
            TensorView::C64($v) => Ok(Tensor::from_typed::<tenferro_tensor::Complex64>($body?)),
        }
    };
}

macro_rules! dispatch_tensor_unary_result {
    ($input:expr, |$tensor:ident| $body:expr) => {{
        let input = $input;
        match input.dtype() {
            DType::F32 => {
                let $tensor = unary_operand::<f32>(input)?;
                Ok(Tensor::from_typed::<f32>($body?))
            }
            DType::F64 => {
                let $tensor = unary_operand::<f64>(input)?;
                Ok(Tensor::from_typed::<f64>($body?))
            }
            DType::I32 => {
                let $tensor = unary_operand::<i32>(input)?;
                Ok(Tensor::from_typed::<i32>($body?))
            }
            DType::I64 => {
                let $tensor = unary_operand::<i64>(input)?;
                Ok(Tensor::from_typed::<i64>($body?))
            }
            DType::Bool => {
                let $tensor = unary_operand::<bool>(input)?;
                Ok(Tensor::from_typed::<bool>($body?))
            }
            DType::C32 => {
                let $tensor = unary_operand::<num_complex::Complex32>(input)?;
                Ok(Tensor::from_typed::<num_complex::Complex32>($body?))
            }
            DType::C64 => {
                let $tensor = unary_operand::<num_complex::Complex64>(input)?;
                Ok(Tensor::from_typed::<num_complex::Complex64>($body?))
            }
            // A caller-owned payload has no CPU operation implementation.
            DType::External(type_id) => Err(crate::Error::unsupported_dtype(
                "cpu",
                tenferro_tensor::DType::External(type_id),
                "an externally defined payload is not supported by this CPU operation",
            )),
        }
    }};
}

/// The typed tensor behind `input`, or the refusal a tag the dispatch cannot resolve reports.
fn unary_operand<T: tenferro_tensor::TensorScalar>(
    input: &Tensor,
) -> crate::Result<&TypedTensor<T>> {
    input.as_typed::<T>().ok_or_else(|| {
        crate::Error::unsupported_dtype(
            "cpu",
            input.dtype(),
            "an externally defined payload is not supported by this CPU operation",
        )
    })
}

macro_rules! dispatch_same_dtype_result {
    ($op:literal, $lhs:expr, $rhs:expr, |$lhs_t:ident, $rhs_t:ident| $body:expr) => {{
        let lhs = $lhs;
        let rhs = $rhs;
        match (lhs.dtype(), rhs.dtype()) {
            (DType::F32, DType::F32) => {
                let ($lhs_t, $rhs_t) = pair_operands::<f32>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<f32>($body?))
            }
            (DType::F64, DType::F64) => {
                let ($lhs_t, $rhs_t) = pair_operands::<f64>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<f64>($body?))
            }
            (DType::I32, DType::I32) => {
                let ($lhs_t, $rhs_t) = pair_operands::<i32>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<i32>($body?))
            }
            (DType::I64, DType::I64) => {
                let ($lhs_t, $rhs_t) = pair_operands::<i64>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<i64>($body?))
            }
            (DType::Bool, DType::Bool) => {
                let ($lhs_t, $rhs_t) = pair_operands::<bool>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<bool>($body?))
            }
            (DType::C32, DType::C32) => {
                let ($lhs_t, $rhs_t) = pair_operands::<num_complex::Complex32>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<num_complex::Complex32>($body?))
            }
            (DType::C64, DType::C64) => {
                let ($lhs_t, $rhs_t) = pair_operands::<num_complex::Complex64>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<num_complex::Complex64>($body?))
            }
            _ => Err(crate::Error::dtype_mismatch($op, lhs.dtype(), rhs.dtype())),
        }
    }};
}

macro_rules! dispatch_same_dtype_without_bool_result {
    ($op:literal, $lhs:expr, $rhs:expr, $bool_message:literal, |$lhs_t:ident, $rhs_t:ident| $body:expr) => {{
        let lhs = $lhs;
        let rhs = $rhs;
        match (lhs.dtype(), rhs.dtype()) {
            (DType::F32, DType::F32) => {
                let ($lhs_t, $rhs_t) = pair_operands::<f32>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<f32>($body?))
            }
            (DType::F64, DType::F64) => {
                let ($lhs_t, $rhs_t) = pair_operands::<f64>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<f64>($body?))
            }
            (DType::I32, DType::I32) => {
                let ($lhs_t, $rhs_t) = pair_operands::<i32>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<i32>($body?))
            }
            (DType::I64, DType::I64) => {
                let ($lhs_t, $rhs_t) = pair_operands::<i64>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<i64>($body?))
            }
            (DType::C32, DType::C32) => {
                let ($lhs_t, $rhs_t) = pair_operands::<num_complex::Complex32>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<num_complex::Complex32>($body?))
            }
            (DType::C64, DType::C64) => {
                let ($lhs_t, $rhs_t) = pair_operands::<num_complex::Complex64>($op, lhs, rhs)?;
                Ok(Tensor::from_typed::<num_complex::Complex64>($body?))
            }
            (DType::Bool, DType::Bool) => Err(crate::Error::unsupported($op, $bool_message)),

            _ => Err(crate::Error::dtype_mismatch($op, lhs.dtype(), rhs.dtype())),
        }
    }};
}

/// Descriptors for a same-dtype pair of reads, passed to `$body` as typed operands.
///
/// The caller has already matched both dtypes to `$ty`, so a mismatch here is a
/// tag/payload disagreement and reports the pair's dtypes.
macro_rules! read_pair_operands {
    ($op:literal, $ty:ty, $variant:ident, $lhs:expr, $rhs:expr, |$a:ident, $b:ident| $body:expr) => {{
        let (lhs, rhs): (&TensorRead<'_>, &TensorRead<'_>) = ($lhs, $rhs);
        let mismatch = || crate::Error::dtype_mismatch($op, lhs.dtype(), rhs.dtype());
        macro_rules! descriptor {
            ($read:expr) => {
                match $read {
                    TensorRead::Tensor(tensor) => {
                        let typed = TensorAsTyped::<$ty>::as_typed(*tensor).ok_or_else(mismatch)?;
                        Operand::from_typed($op, typed)
                    }
                    TensorRead::View(TensorView::$variant(view)) => operand_from_view(view),
                    TensorRead::View(_) => Err(mismatch()),
                }
            };
        }
        let $a = descriptor!(lhs)?;
        let $b = descriptor!(rhs)?;
        Ok(Tensor::from_typed::<$ty>($body?))
    }};
}

/// The typed operands behind a same-dtype pair, or the refusal a mismatched pair reports.
///
/// Callers reach this from a match on the pair of dtypes, so `None` means the tags and the runtime
/// dtypes disagree rather than a caller mistake.
fn pair_operands<'a, T: tenferro_tensor::TensorScalar>(
    op: &'static str,
    lhs: &'a Tensor,
    rhs: &'a Tensor,
) -> crate::Result<(&'a TypedTensor<T>, &'a TypedTensor<T>)> {
    let lhs_t = lhs
        .as_typed::<T>()
        .ok_or_else(|| crate::Error::dtype_mismatch(op, lhs.dtype(), rhs.dtype()))?;
    let rhs_t = rhs
        .as_typed::<T>()
        .ok_or_else(|| crate::Error::dtype_mismatch(op, lhs.dtype(), rhs.dtype()))?;
    Ok((lhs_t, rhs_t))
}

#[cfg(test)]
fn with_test_pool<T>(f: impl FnOnce(&mut BufferPool) -> T) -> T {
    let mut buffers = BufferPool::new();
    f(&mut buffers)
}

#[cfg(test)]
pub(crate) fn gather(
    operand: &Tensor,
    start_indices: &Tensor,
    config: &GatherConfig,
) -> crate::Result<Tensor> {
    with_test_pool(|buffers| {
        let cache = CpuPlanCache::default();
        gather_with_pool(
            buffers,
            &cache,
            operand,
            start_indices,
            config,
            Exec::serial(),
        )
    })
}

pub(crate) fn gather_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: &Tensor,
    start_indices: &Tensor,
    config: &GatherConfig,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let start_indices = try_index_tensor(start_indices)?;
    dispatch_tensor_unary_result!(operand, |t| typed_gather(
        buffers,
        cache,
        t,
        &start_indices,
        config,
        exec,
    ))
}

/// Gather a read, keeping a borrowed operand view's layout instead of staging it.
///
/// The index tensor is converted to a compact `i64` index buffer before the
/// plan runs, so a borrowed index view is materialized here explicitly; the
/// operand view is not copied.
pub(crate) fn gather_read_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: TensorRead<'_>,
    start_indices: TensorRead<'_>,
    config: &GatherConfig,
    mut exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let mut staged_start_indices = None;
    let start_indices = index_tensor_read(
        buffers,
        "gather",
        start_indices,
        &mut staged_start_indices,
        exec.reborrow(),
    )?;
    match operand {
        TensorRead::Tensor(operand) => dispatch_tensor_unary_result!(operand, |t| typed_gather(
            buffers,
            cache,
            t,
            &start_indices,
            config,
            exec,
        )),
        TensorRead::View(view) => {
            dispatch_tensor_view_result!(view, |v| typed_gather_view(
                buffers,
                cache,
                &v,
                &start_indices,
                config,
                exec,
            ))
        }
    }
}

#[cfg(test)]
pub(crate) fn scatter(
    operand: &Tensor,
    scatter_indices: &Tensor,
    updates: &Tensor,
    config: &ScatterConfig,
) -> crate::Result<Tensor> {
    with_test_pool(|buffers| {
        let cache = CpuPlanCache::default();
        scatter_with_pool(
            buffers,
            &cache,
            operand,
            scatter_indices,
            updates,
            config,
            Exec::serial(),
        )
    })
}

pub(crate) fn scatter_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: &Tensor,
    scatter_indices: &Tensor,
    updates: &Tensor,
    config: &ScatterConfig,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let scatter_indices = try_index_tensor(scatter_indices)?;
    dispatch_same_dtype_without_bool_result!(
        "scatter",
        operand,
        updates,
        "Bool data tensors are not supported by additive scatter; supported data dtypes: F32/F64/I32/I64/C32/C64",
        |op, upd| typed_scatter(
            buffers,
            cache,
            op,
            &scatter_indices,
            upd,
            config,
            exec,
        )
    )
}

/// Scatter reads, keeping borrowed operand and update views instead of staging them.
///
/// The index tensor is converted to a compact `i64` index buffer before the
/// plan runs, so a borrowed index view is materialized here explicitly; the
/// operand and update views are not copied.
pub(crate) fn scatter_read_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: TensorRead<'_>,
    scatter_indices: TensorRead<'_>,
    updates: TensorRead<'_>,
    config: &ScatterConfig,
    mut exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let mut staged_scatter_indices = None;
    let scatter_indices = index_tensor_read(
        buffers,
        "scatter",
        scatter_indices,
        &mut staged_scatter_indices,
        exec.reborrow(),
    )?;
    let (operand_dtype, update_dtype) = (operand.dtype(), updates.dtype());
    macro_rules! scatter_pair {
        ($ty:ty, $variant:ident) => {
            read_pair_operands!("scatter", $ty, $variant, &operand, &updates, |op, upd| {
                typed_scatter_at(buffers, cache, op, &scatter_indices, upd, config, exec)
            })
        };
    }
    match (operand_dtype, update_dtype) {
        (DType::F32, DType::F32) => scatter_pair!(f32, F32),
        (DType::F64, DType::F64) => scatter_pair!(f64, F64),
        (DType::I32, DType::I32) => scatter_pair!(i32, I32),
        (DType::I64, DType::I64) => scatter_pair!(i64, I64),
        (DType::C32, DType::C32) => scatter_pair!(tenferro_tensor::Complex32, C32),
        (DType::C64, DType::C64) => scatter_pair!(tenferro_tensor::Complex64, C64),
        (DType::Bool, DType::Bool) => Err(crate::Error::unsupported(
            "scatter",
            "Bool data tensors are not supported by additive scatter; supported data dtypes: F32/F64/I32/I64/C32/C64",
        )),
        _ => Err(crate::Error::dtype_mismatch(
            "scatter",
            operand_dtype,
            update_dtype,
        )),
    }
}

pub(crate) fn try_slice_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    config: &SliceConfig,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    dispatch_tensor_unary_result!(input, |t| typed_slice(buffers, t, config, exec))
}

/// Slice a read, keeping a borrowed view's layout instead of staging it.
pub(crate) fn slice_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    config: &SliceConfig,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => try_slice_with_pool(buffers, input, config, exec),
        TensorRead::View(view) => {
            dispatch_tensor_view_result!(view, |v| typed_slice_view(buffers, &v, config, exec,))
        }
    }
}

#[cfg(test)]
pub(crate) fn dynamic_slice(
    input: &Tensor,
    starts: &Tensor,
    slice_sizes: &[usize],
) -> crate::Result<Tensor> {
    with_test_pool(|buffers| {
        let cache = CpuPlanCache::default();
        dynamic_slice_with_pool(buffers, &cache, input, starts, slice_sizes, Exec::serial())
    })
}

pub(crate) fn dynamic_slice_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    input: &Tensor,
    starts: &Tensor,
    slice_sizes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let starts = try_index_tensor(starts)?;
    dispatch_tensor_unary_result!(input, |t| typed_dynamic_slice(
        buffers,
        cache,
        t,
        &starts,
        slice_sizes,
        exec,
    ))
}

/// Dynamic-slice a read, keeping a borrowed operand view instead of staging it.
///
/// The `starts` vector becomes a compact `i64` index buffer, so a borrowed
/// `starts` view is materialized here explicitly; the operand view is not
/// copied.
pub(crate) fn dynamic_slice_read_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    input: TensorRead<'_>,
    starts: TensorRead<'_>,
    slice_sizes: &[usize],
    mut exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let mut staged_starts = None;
    let starts = index_tensor_read(
        buffers,
        "dynamic_slice",
        starts,
        &mut staged_starts,
        exec.reborrow(),
    )?;
    match input {
        TensorRead::Tensor(input) => dispatch_tensor_unary_result!(input, |t| typed_dynamic_slice(
            buffers,
            cache,
            t,
            &starts,
            slice_sizes,
            exec,
        )),
        TensorRead::View(view) => {
            dispatch_tensor_view_result!(view, |v| typed_dynamic_slice_view(
                buffers,
                cache,
                &v,
                &starts,
                slice_sizes,
                exec,
            ))
        }
    }
}

/// Return `operand` with `update` written at dynamic `starts`.
///
/// Starts are clamped so the whole update window fits inside the operand,
/// matching `dynamic_slice` behavior.
///
/// # Examples
///
/// ```
/// use tenferro_cpu as cpu;
/// use tenferro_tensor::{Tensor, TypedTensor};
///
/// let operand = Tensor::from_typed::<f64>(TypedTensor::from_vec_col_major(vec![5], vec![0.0; 5])?);
/// let update = Tensor::from_typed::<f64>(TypedTensor::from_vec_col_major(vec![2], vec![3.0, 4.0])?);
/// let starts = Tensor::from_typed::<i64>(TypedTensor::from_vec_col_major(vec![1], vec![4])?);
///
/// let out = cpu::dynamic_update_slice(&operand, &update, &starts).unwrap();
/// assert_eq!(out.as_slice::<f64>().unwrap(), &[0.0, 0.0, 0.0, 3.0, 4.0]);
/// # Ok::<(), tenferro_tensor::Error>(())
/// ```
#[cfg(test)]
pub(crate) fn dynamic_update_slice(
    operand: &Tensor,
    update: &Tensor,
    starts: &Tensor,
) -> crate::Result<Tensor> {
    with_test_pool(|buffers| {
        let cache = CpuPlanCache::default();
        dynamic_update_slice_with_pool(buffers, &cache, operand, update, starts, Exec::serial())
    })
}

pub(crate) fn dynamic_update_slice_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: &Tensor,
    update: &Tensor,
    starts: &Tensor,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let starts = try_index_tensor(starts)?;
    dispatch_same_dtype_result!("dynamic_update_slice", operand, update, |op, upd| {
        typed_dynamic_update_slice(buffers, cache, op, upd, &starts, exec)
    })
}

/// Dynamic-update-slice reads, keeping borrowed operand and update views
/// instead of staging them.
///
/// The `starts` vector becomes a compact `i64` index buffer, so a borrowed
/// `starts` view is materialized here explicitly; the operand and update views
/// are not copied.
pub(crate) fn dynamic_update_slice_read_with_pool(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: TensorRead<'_>,
    update: TensorRead<'_>,
    starts: TensorRead<'_>,
    mut exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let mut staged_starts = None;
    let starts = index_tensor_read(
        buffers,
        "dynamic_update_slice",
        starts,
        &mut staged_starts,
        exec.reborrow(),
    )?;
    macro_rules! update_pair {
        ($ty:ty, $variant:ident) => {
            read_pair_operands!(
                "dynamic_update_slice",
                $ty,
                $variant,
                &operand,
                &update,
                |op, upd| typed_dynamic_update_slice_at(buffers, cache, op, upd, &starts, exec,)
            )
        };
    }
    match (operand.dtype(), update.dtype()) {
        (DType::F32, DType::F32) => update_pair!(f32, F32),
        (DType::F64, DType::F64) => update_pair!(f64, F64),
        (DType::I32, DType::I32) => update_pair!(i32, I32),
        (DType::I64, DType::I64) => update_pair!(i64, I64),
        (DType::Bool, DType::Bool) => update_pair!(bool, Bool),
        (DType::C32, DType::C32) => update_pair!(tenferro_tensor::Complex32, C32),
        (DType::C64, DType::C64) => update_pair!(tenferro_tensor::Complex64, C64),
        (lhs, rhs) => Err(crate::Error::dtype_mismatch(
            "dynamic_update_slice",
            lhs,
            rhs,
        )),
    }
}

#[cfg(test)]
pub(crate) fn pad(input: &Tensor, config: &PadConfig) -> crate::Result<Tensor> {
    try_pad(input, config)
}

#[cfg(test)]
fn try_pad(input: &Tensor, config: &PadConfig) -> crate::Result<Tensor> {
    with_test_pool(|buffers| try_pad_with_pool(buffers, input, config, Exec::serial()))
}

pub(crate) fn try_pad_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    config: &PadConfig,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    dispatch_tensor_unary_result!(input, |t| typed_pad(buffers, t, config, exec))
}

/// Pad a read, keeping a borrowed view's layout instead of staging it.
pub(crate) fn pad_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    config: &PadConfig,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => try_pad_with_pool(buffers, input, config, exec),
        TensorRead::View(view) => {
            dispatch_tensor_view_result!(view, |v| typed_pad_view(buffers, &v, config, exec,))
        }
    }
}

/// Concatenate reads, keeping each borrowed view's layout instead of staging it.
pub(crate) fn concatenate_read_with_pool(
    buffers: &mut BufferPool,
    inputs: &[TensorRead<'_>],
    axis: usize,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let first = inputs.first().ok_or_else(|| {
        crate::Error::invalid_argument(
            "concatenate",
            "configuration",
            "concatenate requires at least one input",
        )
    })?;
    let dtype = first.dtype();
    macro_rules! concatenate_operands {
        ($ty:ty, $variant:ident) => {{
            let operands: SmallVec<[Operand<'_, $ty>; 4]> = inputs
                .iter()
                .map(|input| match input {
                    TensorRead::Tensor(tensor) => {
                        let typed = TensorAsTyped::<$ty>::as_typed(*tensor).ok_or_else(|| {
                            crate::Error::dtype_mismatch("concatenate", dtype, tensor.dtype())
                        })?;
                        Operand::from_typed("concatenate", typed)
                    }
                    TensorRead::View(TensorView::$variant(view)) => operand_from_view(view),
                    TensorRead::View(view) => Err(crate::Error::dtype_mismatch(
                        "concatenate",
                        dtype,
                        view.dtype(),
                    )),
                })
                .collect::<crate::Result<_>>()?;
            Ok(Tensor::from_typed::<$ty>(typed_concatenate_operands(
                buffers, &operands, axis, exec,
            )?))
        }};
    }
    match dtype {
        DType::F32 => concatenate_operands!(f32, F32),
        DType::F64 => concatenate_operands!(f64, F64),
        DType::I32 => concatenate_operands!(i32, I32),
        DType::I64 => concatenate_operands!(i64, I64),
        DType::Bool => concatenate_operands!(bool, Bool),
        DType::C32 => concatenate_operands!(tenferro_tensor::Complex32, C32),
        DType::C64 => concatenate_operands!(tenferro_tensor::Complex64, C64),
        // A caller-owned payload has no CPU operation implementation.
        DType::External(type_id) => Err(crate::Error::unsupported_dtype(
            "cpu",
            tenferro_tensor::DType::External(type_id),
            "an externally defined payload is not supported by this CPU operation",
        )),
    }
}

pub(crate) fn try_concatenate_with_pool(
    buffers: &mut BufferPool,
    inputs: &[&Tensor],
    axis: usize,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    let first = inputs.first().copied().ok_or_else(|| {
        crate::Error::invalid_argument(
            "concatenate",
            "configuration",
            "concatenate requires at least one input",
        )
    })?;
    dispatch_tensor_unary_result!(first, |t| typed_concatenate_from_dyn_inputs(
        buffers, t, inputs, axis, exec,
    ))
}

pub(crate) fn reverse_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    axes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    dispatch_tensor_unary_result!(input, |t| typed_reverse(buffers, t, axes, exec))
}

/// Reverse a read, keeping a borrowed view's layout instead of staging it.
pub(crate) fn reverse_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    axes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => reverse_with_pool(buffers, input, axes, exec),
        TensorRead::View(view) => {
            dispatch_tensor_view_result!(view, |v| typed_reverse_view(buffers, &v, axes, exec,))
        }
    }
}

fn typed_slice<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    input: &TypedTensor<T>,
    config: &SliceConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    typed_slice_at(buffers, Operand::from_typed("slice", input)?, config, exec)
}

fn typed_slice_view<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    config: &SliceConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_slice_at(buffers, operand_from_view(view)?, config, exec)
}

fn typed_slice_at<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    operand: Operand<'_, T>,
    config: &SliceConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let Operand {
        data,
        shape: input_shape,
        strides: input_strides,
        offset: input_offset,
    } = operand;
    let rank = input_shape.len();
    if config.starts.len() != rank {
        return Err(crate::Error::rank_mismatch(
            "slice",
            rank,
            config.starts.len(),
        ));
    }
    if config.limits.len() != rank {
        return Err(crate::Error::rank_mismatch(
            "slice",
            rank,
            config.limits.len(),
        ));
    }
    if config.strides.len() != rank {
        return Err(crate::Error::rank_mismatch(
            "slice",
            rank,
            config.strides.len(),
        ));
    }

    let out_shape: Vec<usize> = input_shape
        .iter()
        .enumerate()
        .map(|(axis, &dim)| {
            let start = config.starts[axis];
            let limit = config.limits[axis];
            let stride = config.strides[axis];
            if start > limit {
                return Err(crate::Error::invalid_argument(
                    "slice",
                    "configuration",
                    format!("start exceeds limit on axis {axis}"),
                ));
            }
            // INVARIANT: This boundary check intentionally mirrors CUDA's
            // validator. CPU and GPU are independent backend leaves, and
            // sharing it via tenferro-tensor would require a new public
            // validation API.
            if limit > dim {
                return Err(crate::Error::invalid_argument(
                    "slice",
                    "configuration",
                    format!("limit {limit} on axis {axis} exceeds dimension size {dim}"),
                ));
            }
            if stride == 0 {
                return Err(crate::Error::invalid_argument(
                    "slice",
                    "configuration",
                    format!("stride must be positive on axis {axis}"),
                ));
            }
            let span = limit - start;
            Ok(span.div_ceil(stride))
        })
        .collect::<crate::Result<Vec<_>>>()?;

    let dtype = kernel_dtype(T::dtype());
    let out_strides = inline_col_major_strides("slice", &out_shape)?;
    let plan = ErasedSlicePlan::compile(
        dtype,
        input_shape,
        input_strides,
        &out_shape,
        &out_strides,
        &config.starts,
        &config.limits,
        &config.strides,
    )
    .map_err(|err| crate::Error::backend_source("slice", err))?;

    let mut out = PooledUninitOutput::new(buffers, &out_shape)?;
    let input_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(data),
            input_shape,
            input_strides,
            input_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("slice", err))?;
    let input_ptr = ErasedRawStridedPtr::from_ref(&input_ref);
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            &out_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("slice", err))?;
    plan.execute_uninit(&mut out_ref, &input_ptr, exec)
        .map_err(|err| crate::Error::backend_source("slice", err))?;

    // INVARIANT: the compact output has no unreachable storage, so successful
    // static-slice replay initializes every element.
    // SAFETY: every element contains a valid T after successful replay.
    unsafe { out.assume_init_recycled() }
}

fn typed_concatenate_from_dyn_inputs<T>(
    buffers: &mut BufferPool,
    _first: &TypedTensor<T>,
    inputs: &[&Tensor],
    axis: usize,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar,
    Tensor: TensorAsTyped<T>,
{
    let first_dtype = inputs
        .first()
        .ok_or_else(|| {
            crate::Error::invalid_argument(
                "concatenate",
                "configuration",
                "concatenate requires at least one input",
            )
        })?
        .dtype();
    let typed_inputs = collect_typed_inputs(first_dtype, inputs)?;
    typed_concatenate(buffers, &typed_inputs, axis, exec)
}

fn collect_typed_inputs<'a, T>(
    first_dtype: crate::DType,
    inputs: &[&'a Tensor],
) -> crate::Result<SmallVec<[&'a TypedTensor<T>; 4]>>
where
    Tensor: TensorAsTyped<T>,
{
    inputs
        .iter()
        .map(|tensor| {
            TensorAsTyped::<T>::as_typed(*tensor).ok_or_else(|| {
                crate::Error::dtype_mismatch("concatenate", first_dtype, tensor.dtype())
            })
        })
        .collect()
}

fn typed_concatenate<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    inputs: &[&TypedTensor<T>],
    axis: usize,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let operands: SmallVec<[Operand<'_, T>; 4]> = inputs
        .iter()
        .map(|input| Operand::from_typed("concatenate", input))
        .collect::<crate::Result<_>>()?;
    typed_concatenate_operands(buffers, &operands, axis, exec)
}

fn typed_concatenate_operands<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    inputs: &[Operand<'_, T>],
    axis: usize,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let first = inputs.first().ok_or_else(|| {
        crate::Error::invalid_argument(
            "concatenate",
            "configuration",
            "concatenate requires at least one input",
        )
    })?;
    let first_shape = first.shape;
    let rank = first_shape.len();
    if axis >= rank {
        return Err(crate::Error::axis_out_of_bounds("concatenate", axis, rank));
    }

    let mut out_shape = tenferro_tensor::ShapeVec::from(first_shape);
    let mut axis_extent = 0usize;
    for input in inputs {
        let input_shape = input.shape;
        if input_shape.len() != rank {
            return Err(crate::Error::rank_mismatch(
                "concatenate",
                rank,
                input_shape.len(),
            ));
        }
        for dim in 0..rank {
            if dim == axis {
                axis_extent = axis_extent.checked_add(input_shape[dim]).ok_or_else(|| {
                    crate::Error::invalid_argument(
                        "concatenate",
                        "configuration",
                        "concatenate axis extent overflows usize".to_string(),
                    )
                })?;
            } else if input_shape[dim] != first_shape[dim] {
                return Err(crate::Error::shape_mismatch(
                    "concatenate",
                    first_shape.to_vec(),
                    input_shape.to_vec(),
                ));
            }
        }
    }
    out_shape[axis] = axis_extent;

    let dtype = kernel_dtype(T::dtype());
    let input_dims: TinyVec<[&[usize]; 4]> = inputs.iter().map(|input| input.shape).collect();
    let input_stride_refs: TinyVec<[&[isize]; 4]> =
        inputs.iter().map(|input| input.strides).collect();
    let out_strides = inline_col_major_strides("concatenate", &out_shape)?;
    let plan = ErasedConcatenatePlan::compile(
        dtype,
        &input_dims,
        &input_stride_refs,
        &out_shape,
        &out_strides,
        axis,
    )
    .map_err(|err| crate::Error::backend_source("concatenate", err))?;

    let mut out = PooledUninitOutput::new(buffers, &out_shape)?;
    // `ErasedRawStridedRef` has no `Default`, so `tinyvec::Array` cannot hold
    // it; `SmallVec` keeps the four-operand descriptor list inline.
    let input_refs: SmallVec<[ErasedRawStridedRef<'_>; 4]> = inputs
        .iter()
        .map(|input| {
            unsafe {
                // SAFETY: this operation supplies initialized typed bytes with
                // alignment and dtype matching the descriptor; validated dimensions,
                // strides, and offset keep every reachable source element in bounds
                // for the retained borrow.
                crate::erased_raw_strided_ref(
                    dtype,
                    typed_bytes(input.data),
                    input.shape,
                    input.strides,
                    input.offset,
                )
            }
            .map_err(|err| crate::Error::backend_source("concatenate", err))
        })
        .collect::<crate::Result<_>>()?;
    // `ErasedRawStridedPtr` has no `Default`; `SmallVec` keeps the short list
    // inline as above.
    let input_ptrs: SmallVec<[ErasedRawStridedPtr<'_>; 4]> = input_refs
        .iter()
        .map(ErasedRawStridedPtr::from_ref)
        .collect();
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            &out_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("concatenate", err))?;
    plan.execute_uninit(&mut out_ref, &input_ptrs, exec)
        .map_err(|err| crate::Error::backend_source("concatenate", err))?;

    // INVARIANT: the compact output has no unreachable storage, so successful
    // concatenate replay initializes every element.
    // SAFETY: every element contains a valid T after successful replay.
    unsafe { out.assume_init_recycled() }
}

fn typed_reverse<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    input: &TypedTensor<T>,
    axes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    typed_reverse_at(buffers, Operand::from_typed("reverse", input)?, axes, exec)
}

fn typed_reverse_view<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    axes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_reverse_at(buffers, operand_from_view(view)?, axes, exec)
}

fn typed_reverse_at<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    operand: Operand<'_, T>,
    axes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let Operand {
        data,
        shape: input_shape,
        strides: input_strides,
        offset: input_offset,
    } = operand;
    let rank = input_shape.len();
    for &axis in axes {
        if axis >= rank {
            return Err(crate::Error::axis_out_of_bounds("reverse", axis, rank));
        }
    }

    let dtype = kernel_dtype(T::dtype());
    let out_strides = inline_col_major_strides("reverse", input_shape)?;
    let plan = ErasedReversePlan::compile(dtype, input_shape, input_strides, &out_strides, axes)
        .map_err(|err| crate::Error::backend_source("reverse", err))?;

    let mut out = PooledUninitOutput::new(buffers, input_shape)?;
    let input_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(data),
            input_shape,
            input_strides,
            input_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("reverse", err))?;
    let input_ptr = ErasedRawStridedPtr::from_ref(&input_ref);
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            input_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("reverse", err))?;
    plan.execute_uninit(&mut out_ref, &input_ptr, exec)
        .map_err(|err| crate::Error::backend_source("reverse", err))?;

    // INVARIANT: the compact output has no unreachable storage, so successful
    // reverse replay initializes every element.
    // SAFETY: every element contains a valid T after successful replay.
    unsafe { out.assume_init_recycled() }
}

/// Index values for one gather, scatter, or dynamic-slice call.
///
/// An `I64` index tensor is borrowed in place; every other supported index dtype
/// is converted once into an owned `i64` buffer because the kernel reads `i64`.
enum IndexValues<'a> {
    Borrowed(&'a [i64]),
    Converted(Vec<i64>),
}

impl IndexValues<'_> {
    fn as_slice(&self) -> &[i64] {
        match self {
            Self::Borrowed(values) => values,
            Self::Converted(values) => values,
        }
    }
}

struct IndexTensor<'a> {
    shape: Vec<usize>,
    values: IndexValues<'a>,
}

/// Maximum exact integer representable by f32 (2^24).
const F32_MAX_EXACT_INT: f32 = 16_777_216.0;
/// Maximum exact integer representable by f64 (2^53).
const F64_MAX_EXACT_INT: f64 = 9_007_199_254_740_992.0;

fn f32_index_to_i64(value: f32) -> crate::Result<i64> {
    if !value.is_finite() || value.fract() != 0.0 || value.abs() > F32_MAX_EXACT_INT {
        return Err(crate::Error::invalid_argument(
            "index_tensor",
            "index",
            format!("index value {value} is not an exactly representable i64"),
        ));
    }
    Ok(value as i64)
}

fn f64_index_to_i64(value: f64) -> crate::Result<i64> {
    if !value.is_finite() || value.fract() != 0.0 || value.abs() > F64_MAX_EXACT_INT {
        return Err(crate::Error::invalid_argument(
            "index_tensor",
            "index",
            format!("index value {value} is not an exactly representable i64"),
        ));
    }
    Ok(value as i64)
}

/// The refusal this module produces for a dtype it cannot read as an index tensor.
///
/// The accessor returns it when the tag table and the runtime dtype disagree, which
/// is unreachable in practice; the tags that are genuinely rejected keep their own
/// arms and the same expression.
fn unsupported_index_dtype(complex: bool) -> crate::Error {
    crate::Error::invalid_argument(
        "index_tensor",
        "configuration",
        if complex {
            "complex index tensors are not supported; supported index dtypes: I32/I64/F32/F64"
        } else {
            "bool index tensors are not supported; supported index dtypes: I32/I64/F32/F64"
        },
    )
}

/// Index values from a read: an owned tensor is used in place, and only a view
/// is materialized (into `staged`), because the plans read compact indices.
fn index_tensor_read<'a>(
    buffers: &mut BufferPool,
    op: &'static str,
    read: TensorRead<'a>,
    staged: &'a mut Option<Tensor>,
    exec: Exec<'_>,
) -> crate::Result<IndexTensor<'a>> {
    let tensor = match read {
        TensorRead::Tensor(tensor) => tensor,
        view => staged.insert(crate::materialize_tensor_read(buffers, op, view, exec)?),
    };
    try_index_tensor(tensor)
}

fn try_index_tensor(tensor: &Tensor) -> crate::Result<IndexTensor<'_>> {
    match tensor.dtype() {
        DType::I32 => {
            let t = tensor
                .as_typed::<i32>()
                .ok_or_else(|| unsupported_index_dtype(false))?;
            Ok(IndexTensor {
                shape: t.shape().to_vec(),
                values: IndexValues::Converted(
                    typed_host_data("index_tensor", t)?
                        .iter()
                        .map(|&value| value as i64)
                        .collect(),
                ),
            })
        }
        DType::I64 => {
            let t = tensor
                .as_typed::<i64>()
                .ok_or_else(|| unsupported_index_dtype(false))?;
            // Borrow the index storage instead of copying it per call; the
            // kernel only reads `i64` and the caller keeps the tensor alive.
            Ok(IndexTensor {
                shape: t.shape().to_vec(),
                values: IndexValues::Borrowed(typed_host_data("index_tensor", t)?),
            })
        }
        DType::F32 => {
            let t = tensor
                .as_typed::<f32>()
                .ok_or_else(|| unsupported_index_dtype(false))?;
            let values: crate::Result<Vec<i64>> = typed_host_data("index_tensor", t)?
                .iter()
                .map(|&value| f32_index_to_i64(value))
                .collect();
            Ok(IndexTensor {
                shape: t.shape().to_vec(),
                values: IndexValues::Converted(values?),
            })
        }
        DType::F64 => {
            let t = tensor
                .as_typed::<f64>()
                .ok_or_else(|| unsupported_index_dtype(false))?;
            let values: crate::Result<Vec<i64>> = typed_host_data("index_tensor", t)?
                .iter()
                .map(|&value| f64_index_to_i64(value))
                .collect();
            Ok(IndexTensor {
                shape: t.shape().to_vec(),
                values: IndexValues::Converted(values?),
            })
        }
        DType::Bool => Err(unsupported_index_dtype(false)),
        DType::C32 | DType::C64 | DType::External(_) => Err(unsupported_index_dtype(true)),
    }
}

fn checked_product(op: &'static str, role: &'static str, shape: &[usize]) -> crate::Result<usize> {
    shape.iter().try_fold(1usize, |acc, &dim| {
        acc.checked_mul(dim).ok_or_else(|| {
            crate::Error::invalid_argument(
                op,
                "configuration",
                format!("{role} element count overflows usize"),
            )
        })
    })
}

fn try_index_vector_size(
    op: &'static str,
    shape: &[usize],
    index_vector_dim: usize,
) -> crate::Result<usize> {
    if index_vector_dim > shape.len() {
        return Err(crate::Error::axis_out_of_bounds(
            op,
            index_vector_dim,
            shape.len(),
        ));
    }
    Ok(if index_vector_dim == shape.len() {
        1
    } else {
        shape[index_vector_dim]
    })
}

fn try_index_batch_shape(
    op: &'static str,
    shape: &[usize],
    index_vector_dim: usize,
) -> crate::Result<Vec<usize>> {
    if index_vector_dim > shape.len() {
        return Err(crate::Error::axis_out_of_bounds(
            op,
            index_vector_dim,
            shape.len(),
        ));
    }
    if index_vector_dim == shape.len() {
        return Ok(shape.to_vec());
    }
    Ok(shape
        .iter()
        .enumerate()
        .filter_map(|(axis, &dim)| (axis != index_vector_dim).then_some(dim))
        .collect())
}

fn clamp_window_start(
    op: &'static str,
    start: i64,
    dim_size: usize,
    window_size: usize,
) -> crate::Result<usize> {
    if window_size > dim_size {
        return Err(crate::Error::invalid_argument(
            op,
            "configuration",
            format!("window size {window_size} exceeds dimension size {dim_size}"),
        ));
    }
    // An empty axis can be longer than `i64::MAX`, so the bound saturates
    // rather than wrapping negative.
    let max_start = i64::try_from(dim_size - window_size).unwrap_or(i64::MAX);
    Ok(start.clamp(0, max_start) as usize)
}

fn operand_window_dims(rank: usize, collapsed_or_inserted: &[usize]) -> Vec<usize> {
    (0..rank)
        .filter(|dim| !collapsed_or_inserted.contains(dim))
        .collect()
}

fn typed_gather<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: &TypedTensor<T>,
    start_indices: &IndexTensor,
    config: &GatherConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    typed_gather_at(
        buffers,
        cache,
        Operand::from_typed("gather", operand)?,
        start_indices,
        config,
        exec,
    )
}

fn typed_gather_view<T, R>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: &TypedTensorView<'_, T, R>,
    start_indices: &IndexTensor,
    config: &GatherConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_gather_at(
        buffers,
        cache,
        operand_from_view(operand)?,
        start_indices,
        config,
        exec,
    )
}

fn typed_gather_at<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: Operand<'_, T>,
    start_indices: &IndexTensor,
    config: &GatherConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let Operand {
        data,
        shape: operand_shape,
        strides: operand_strides,
        offset: operand_offset,
    } = operand;
    let rank = operand_shape.len();
    if config.slice_sizes.len() != rank {
        return Err(crate::Error::rank_mismatch(
            "gather",
            rank,
            config.slice_sizes.len(),
        ));
    }

    for &dim in &config.collapsed_slice_dims {
        if dim >= rank {
            return Err(crate::Error::axis_out_of_bounds("gather", dim, rank));
        }
    }
    {
        let mut seen = vec![false; rank];
        for &dim in &config.collapsed_slice_dims {
            if seen[dim] {
                return Err(crate::Error::duplicate_axis(
                    "gather",
                    dim,
                    "collapsed_slice_dims",
                ));
            }
            seen[dim] = true;
        }
    }
    for &dim in &config.collapsed_slice_dims {
        if config.slice_sizes[dim] != 1 {
            return Err(crate::Error::invalid_argument(
                "gather",
                "configuration",
                format!(
                    "collapsed slice dimension {dim} must have slice_size == 1, got {}",
                    config.slice_sizes[dim]
                ),
            ));
        }
    }

    let index_size =
        try_index_vector_size("gather", &start_indices.shape, config.index_vector_dim)?;
    if index_size != config.start_index_map.len() {
        return Err(crate::Error::invalid_argument(
            "gather",
            "configuration",
            format!(
                "start_index_map length {} does not match index vector size {}",
                config.start_index_map.len(),
                index_size
            ),
        ));
    }
    for &operand_dim in &config.start_index_map {
        if operand_dim >= rank {
            return Err(crate::Error::axis_out_of_bounds(
                "gather",
                operand_dim,
                rank,
            ));
        }
    }
    {
        let mut seen = vec![false; rank];
        for &operand_dim in &config.start_index_map {
            if seen[operand_dim] {
                return Err(crate::Error::duplicate_axis(
                    "gather",
                    operand_dim,
                    "start_index_map",
                ));
            }
            seen[operand_dim] = true;
        }
    }

    let window_dims = operand_window_dims(rank, &config.collapsed_slice_dims);
    if config.offset_dims.len() != window_dims.len() {
        return Err(crate::Error::invalid_argument(
            "gather",
            "configuration",
            format!(
                "offset_dims length {} does not match window dims count {}",
                config.offset_dims.len(),
                window_dims.len()
            ),
        ));
    }

    let batch_shape =
        try_index_batch_shape("gather", &start_indices.shape, config.index_vector_dim)?;
    let out_rank = batch_shape.len() + config.offset_dims.len();
    for &out_axis in &config.offset_dims {
        if out_axis >= out_rank {
            return Err(crate::Error::axis_out_of_bounds(
                "gather", out_axis, out_rank,
            ));
        }
    }
    {
        let mut seen = vec![false; out_rank];
        for &out_axis in &config.offset_dims {
            if seen[out_axis] {
                return Err(crate::Error::duplicate_axis(
                    "gather",
                    out_axis,
                    "offset_dims",
                ));
            }
            seen[out_axis] = true;
        }
    }

    let mut out_axis_to_operand_dim = vec![None; out_rank];
    for (offset_axis, &out_axis) in config.offset_dims.iter().enumerate() {
        out_axis_to_operand_dim[out_axis] = Some(window_dims[offset_axis]);
    }

    let mut out_shape = vec![0usize; out_rank];
    let mut batch_axis = 0usize;
    for out_axis in 0..out_rank {
        if let Some(operand_dim) = out_axis_to_operand_dim[out_axis] {
            out_shape[out_axis] = config.slice_sizes[operand_dim];
        } else {
            out_shape[out_axis] = batch_shape[batch_axis];
            batch_axis += 1;
        }
    }

    for (operand_dim, &dim_size) in operand_shape.iter().enumerate() {
        let _ = clamp_window_start("gather", 0, dim_size, config.slice_sizes[operand_dim])?;
    }

    let mut out = PooledUninitOutput::<T>::new(buffers, &out_shape)?;

    let dtype = kernel_dtype(T::dtype());
    let index_strides = col_major_strides(&start_indices.shape);
    let out_strides = col_major_strides(&out_shape);
    let index_dtype = KernelDType::I64;
    let key = indexed_plan_key(
        IndexedPlanFamily::Gather,
        dtype,
        index_dtype,
        &[operand_shape, &start_indices.shape, &out_shape],
        &[operand_strides, &index_strides, &out_strides],
        &[
            &config.offset_dims,
            &config.collapsed_slice_dims,
            &config.start_index_map,
            std::slice::from_ref(&config.index_vector_dim),
            &config.slice_sizes,
        ],
    );
    let plan = cache
        .gather(key, || {
            ErasedGatherPlan::compile(
                dtype,
                index_dtype,
                operand_shape,
                operand_strides,
                &start_indices.shape,
                &index_strides,
                &out_shape,
                &out_strides,
                gather_spec(config),
            )
        })
        .map_err(|err| crate::Error::backend_source("gather", err))?;

    let operand_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(data),
            operand_shape,
            operand_strides,
            operand_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("gather", err))?;
    let index_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            index_dtype,
            typed_bytes(start_indices.values.as_slice()),
            &start_indices.shape,
            &index_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("gather", err))?;
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            &out_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("gather", err))?;
    let operand_ptr = ErasedRawStridedPtr::from_ref(&operand_ref);
    let index_ptr = ErasedRawStridedPtr::from_ref(&index_ref);
    plan.execute_uninit(&mut out_ref, &operand_ptr, &index_ptr, exec)
        .map_err(|err| crate::Error::backend_source("gather", err))?;

    // SAFETY: the gather plan writes every logical output element.
    unsafe { out.assume_init_recycled() }
}

fn typed_scatter<T>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: &TypedTensor<T>,
    scatter_indices: &IndexTensor,
    updates: &TypedTensor<T>,
    config: &ScatterConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar,
{
    typed_scatter_at(
        buffers,
        cache,
        Operand::from_typed("scatter", operand)?,
        scatter_indices,
        Operand::from_typed("scatter", updates)?,
        config,
        exec,
    )
}

fn typed_scatter_at<T>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: Operand<'_, T>,
    scatter_indices: &IndexTensor,
    updates: Operand<'_, T>,
    config: &ScatterConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar,
{
    let Operand {
        data: operand_data,
        shape: operand_shape,
        strides: operand_strides,
        offset: operand_offset,
    } = operand;
    let Operand {
        data: update_data,
        shape: updates_shape,
        strides: update_strides,
        offset: update_offset,
    } = updates;
    let op_rank = operand_shape.len();
    for &dim in &config.inserted_window_dims {
        if dim >= op_rank {
            return Err(crate::Error::axis_out_of_bounds("scatter", dim, op_rank));
        }
    }
    {
        let mut seen = vec![false; op_rank];
        for &dim in &config.inserted_window_dims {
            if seen[dim] {
                return Err(crate::Error::duplicate_axis(
                    "scatter",
                    dim,
                    "inserted_window_dims",
                ));
            }
            seen[dim] = true;
        }
    }

    let index_size =
        try_index_vector_size("scatter", &scatter_indices.shape, config.index_vector_dim)?;
    if index_size != config.scatter_dims_to_operand_dims.len() {
        return Err(crate::Error::invalid_argument(
            "scatter",
            "configuration",
            format!(
                "scatter_dims_to_operand_dims length {} does not match index vector size {}",
                config.scatter_dims_to_operand_dims.len(),
                index_size
            ),
        ));
    }
    for &operand_dim in &config.scatter_dims_to_operand_dims {
        if operand_dim >= op_rank {
            return Err(crate::Error::axis_out_of_bounds(
                "scatter",
                operand_dim,
                op_rank,
            ));
        }
    }
    {
        let mut seen = vec![false; op_rank];
        for &operand_dim in &config.scatter_dims_to_operand_dims {
            if seen[operand_dim] {
                return Err(crate::Error::duplicate_axis(
                    "scatter",
                    operand_dim,
                    "scatter_dims_to_operand_dims",
                ));
            }
            seen[operand_dim] = true;
        }
    }

    let batch_shape =
        try_index_batch_shape("scatter", &scatter_indices.shape, config.index_vector_dim)?;
    let window_dims = operand_window_dims(op_rank, &config.inserted_window_dims);
    if config.update_window_dims.len() != window_dims.len() {
        return Err(crate::Error::invalid_argument(
            "scatter",
            "configuration",
            format!(
                "update_window_dims length {} does not match window dims count {}",
                config.update_window_dims.len(),
                window_dims.len()
            ),
        ));
    }

    let update_rank = updates_shape.len();
    let expected_batch_rank = update_rank
        .checked_sub(config.update_window_dims.len())
        .ok_or_else(|| {
            crate::Error::invalid_argument(
                "scatter",
                "configuration",
                format!(
                    "update_window_dims length {} exceeds update rank {}",
                    config.update_window_dims.len(),
                    update_rank
                ),
            )
        })?;
    if expected_batch_rank != batch_shape.len() {
        return Err(crate::Error::invalid_argument(
            "scatter",
            "configuration",
            format!(
                "updates batch rank {} does not match index batch shape length {}",
                expected_batch_rank,
                batch_shape.len()
            ),
        ));
    }

    for &axis in &config.update_window_dims {
        if axis >= update_rank {
            return Err(crate::Error::axis_out_of_bounds(
                "scatter",
                axis,
                update_rank,
            ));
        }
    }
    {
        let mut seen = vec![false; update_rank];
        for &axis in &config.update_window_dims {
            if seen[axis] {
                return Err(crate::Error::duplicate_axis(
                    "scatter",
                    axis,
                    "update_window_dims",
                ));
            }
            seen[axis] = true;
        }
    }

    let mut is_update_window_dim = vec![false; update_rank];
    for &axis in &config.update_window_dims {
        is_update_window_dim[axis] = true;
    }

    {
        let mut batch_axis = 0usize;
        for axis in 0..update_rank {
            if !is_update_window_dim[axis] {
                if updates_shape[axis] != batch_shape[batch_axis] {
                    return Err(crate::Error::invalid_argument("scatter", "configuration", format!(
                            "updates batch dim {} extent {} does not match index batch dim {} extent {}",
                            axis,
                            updates_shape[axis],
                            batch_axis,
                            batch_shape[batch_axis]
                        )));
                }
                batch_axis += 1;
            }
        }
    }

    let mut window_shape = vec![1usize; op_rank];
    let mut window_shape_updates = vec![0usize; config.update_window_dims.len()];
    for (pos, &update_axis) in config.update_window_dims.iter().enumerate() {
        let dim = updates_shape[update_axis];
        window_shape_updates[pos] = dim;
        window_shape[window_dims[pos]] = dim;
    }
    for axis in 0..op_rank {
        let _ = clamp_window_start("scatter", 0, operand_shape[axis], window_shape[axis])?;
    }

    checked_product("scatter", "batch shape", &batch_shape)?;
    checked_product("scatter", "window update shape", &window_shape_updates)?;

    let dtype = kernel_dtype(T::dtype());
    let index_dtype = KernelDType::I64;
    let index_strides = col_major_strides(&scatter_indices.shape);
    let out_strides = col_major_strides(operand_shape);
    let key = indexed_plan_key(
        IndexedPlanFamily::Scatter,
        dtype,
        index_dtype,
        &[
            operand_shape,
            &scatter_indices.shape,
            updates_shape,
            operand_shape,
        ],
        &[
            operand_strides,
            &index_strides,
            update_strides,
            &out_strides,
        ],
        &[
            &config.update_window_dims,
            &config.inserted_window_dims,
            &config.scatter_dims_to_operand_dims,
            std::slice::from_ref(&config.index_vector_dim),
        ],
    );
    let plan = cache
        .scatter(key, || {
            ErasedScatterPlan::compile(
                dtype,
                index_dtype,
                operand_shape,
                operand_strides,
                &scatter_indices.shape,
                &index_strides,
                updates_shape,
                update_strides,
                operand_shape,
                &out_strides,
                scatter_spec(config),
            )
        })
        .map_err(|err| crate::Error::backend_source("scatter", err))?;

    // INVARIANT: ErasedScatterPlan first copies the full operand into `out`,
    // then applies every additive update.
    let mut out = PooledUninitOutput::<T>::new(buffers, operand_shape)?;
    let operand_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(operand_data),
            operand_shape,
            operand_strides,
            operand_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("scatter", err))?;
    let index_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            index_dtype,
            typed_bytes(scatter_indices.values.as_slice()),
            &scatter_indices.shape,
            &index_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("scatter", err))?;
    let update_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(update_data),
            updates_shape,
            update_strides,
            update_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("scatter", err))?;
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            operand_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("scatter", err))?;
    let operand_ptr = ErasedRawStridedPtr::from_ref(&operand_ref);
    let index_ptr = ErasedRawStridedPtr::from_ref(&index_ref);
    let update_ptr = ErasedRawStridedPtr::from_ref(&update_ref);
    plan.execute_uninit(&mut out_ref, &operand_ptr, &index_ptr, &update_ptr, exec)
        .map_err(|err| crate::Error::backend_source("scatter", err))?;

    // SAFETY: the scatter plan first copies the initialized operand and then
    // applies additive updates before returning success.
    unsafe { out.assume_init_recycled() }
}

fn typed_dynamic_slice<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    input: &TypedTensor<T>,
    starts: &IndexTensor,
    slice_sizes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    typed_dynamic_slice_at(
        buffers,
        cache,
        Operand::from_typed("dynamic_slice", input)?,
        starts,
        slice_sizes,
        exec,
    )
}

fn typed_dynamic_slice_view<T, R>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    input: &TypedTensorView<'_, T, R>,
    starts: &IndexTensor,
    slice_sizes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_dynamic_slice_at(
        buffers,
        cache,
        operand_from_view(input)?,
        starts,
        slice_sizes,
        exec,
    )
}

fn typed_dynamic_slice_at<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    input: Operand<'_, T>,
    starts: &IndexTensor,
    slice_sizes: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let Operand {
        data,
        shape: input_shape,
        strides: input_strides,
        offset: input_offset,
    } = input;
    if slice_sizes.len() != input_shape.len() {
        return Err(crate::Error::rank_mismatch(
            "dynamic_slice",
            input_shape.len(),
            slice_sizes.len(),
        ));
    }
    if starts.shape.len() != 1 {
        return Err(crate::Error::invalid_argument(
            "dynamic_slice",
            "starts",
            "starts must be a rank-1 tensor",
        ));
    }
    if starts.values.as_slice().len() != input_shape.len() {
        return Err(crate::Error::invalid_argument(
            "dynamic_slice",
            "starts",
            format!(
                "starts length {} must match input rank {}",
                starts.values.as_slice().len(),
                input_shape.len()
            ),
        ));
    }

    for axis in 0..input_shape.len() {
        let _ = clamp_window_start(
            "dynamic_slice",
            starts.values.as_slice()[axis],
            input_shape[axis],
            slice_sizes[axis],
        )?;
    }

    let out_shape = slice_sizes.to_vec();
    let dtype = kernel_dtype(T::dtype());
    let start_strides = col_major_strides(&starts.shape);
    let out_strides = col_major_strides(&out_shape);
    let index_dtype = KernelDType::I64;
    let key = indexed_plan_key(
        IndexedPlanFamily::DynamicSlice,
        dtype,
        index_dtype,
        &[input_shape, &starts.shape, &out_shape],
        &[input_strides, &start_strides, &out_strides],
        &[slice_sizes],
    );
    let plan = cache
        .dynamic_slice(key, || {
            ErasedDynamicSlicePlan::compile(
                dtype,
                index_dtype,
                input_shape,
                input_strides,
                &starts.shape,
                &start_strides,
                &out_shape,
                &out_strides,
                slice_sizes,
            )
        })
        .map_err(|err| crate::Error::backend_source("dynamic_slice", err))?;
    // INVARIANT: ErasedDynamicSlicePlan writes every output coordinate exactly once.
    let mut out = PooledUninitOutput::<T>::new(buffers, &out_shape)?;
    let input_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(data),
            input_shape,
            input_strides,
            input_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("dynamic_slice", err))?;
    let start_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            index_dtype,
            typed_bytes(starts.values.as_slice()),
            &starts.shape,
            &start_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("dynamic_slice", err))?;
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            &out_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("dynamic_slice", err))?;
    let input_ptr = ErasedRawStridedPtr::from_ref(&input_ref);
    let start_ptr = ErasedRawStridedPtr::from_ref(&start_ref);
    plan.execute_uninit(&mut out_ref, &input_ptr, &start_ptr, exec)
        .map_err(|err| crate::Error::backend_source("dynamic_slice", err))?;

    // SAFETY: the dynamic-slice plan writes every logical output element.
    unsafe { out.assume_init_recycled() }
}

fn typed_dynamic_update_slice<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: &TypedTensor<T>,
    update: &TypedTensor<T>,
    starts: &IndexTensor,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    typed_dynamic_update_slice_at(
        buffers,
        cache,
        Operand::from_typed("dynamic_update_slice", operand)?,
        Operand::from_typed("dynamic_update_slice", update)?,
        starts,
        exec,
    )
}

fn typed_dynamic_update_slice_at<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    cache: &CpuPlanCache,
    operand: Operand<'_, T>,
    update: Operand<'_, T>,
    starts: &IndexTensor,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let Operand {
        data: operand_data,
        shape: operand_shape,
        strides: operand_strides,
        offset: operand_offset,
    } = operand;
    let Operand {
        data: update_data,
        shape: update_shape,
        strides: update_strides,
        offset: update_offset,
    } = update;
    if update_shape.len() != operand_shape.len() {
        return Err(crate::Error::rank_mismatch(
            "dynamic_update_slice",
            operand_shape.len(),
            update_shape.len(),
        ));
    }
    if starts.shape.len() != 1 {
        return Err(crate::Error::invalid_argument(
            "dynamic_update_slice",
            "configuration",
            "starts must be a rank-1 tensor",
        ));
    }
    if starts.values.as_slice().len() != operand_shape.len() {
        return Err(crate::Error::invalid_argument(
            "dynamic_update_slice",
            "configuration",
            format!(
                "starts length {} must match operand rank {}",
                starts.values.as_slice().len(),
                operand_shape.len()
            ),
        ));
    }

    for axis in 0..operand_shape.len() {
        let _ = clamp_window_start(
            "dynamic_update_slice",
            starts.values.as_slice()[axis],
            operand_shape[axis],
            update_shape[axis],
        )?;
    }

    let dtype = kernel_dtype(T::dtype());
    let start_strides = col_major_strides(&starts.shape);
    let out_strides = col_major_strides(operand_shape);
    let index_dtype = KernelDType::I64;
    let key = indexed_plan_key(
        IndexedPlanFamily::DynamicUpdateSlice,
        dtype,
        index_dtype,
        &[operand_shape, &starts.shape, update_shape, operand_shape],
        &[
            operand_strides,
            &start_strides,
            update_strides,
            &out_strides,
        ],
        &[],
    );
    let plan = cache
        .dynamic_update_slice(key, || {
            ErasedDynamicUpdateSlicePlan::compile(
                dtype,
                index_dtype,
                operand_shape,
                operand_strides,
                &starts.shape,
                &start_strides,
                update_shape,
                update_strides,
                operand_shape,
                &out_strides,
            )
        })
        .map_err(|err| crate::Error::backend_source("dynamic_update_slice", err))?;
    // INVARIANT: ErasedDynamicUpdateSlicePlan copies the full operand into
    // `out` before overwriting the update window.
    let mut out = PooledUninitOutput::<T>::new(buffers, operand_shape)?;
    let operand_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(operand_data),
            operand_shape,
            operand_strides,
            operand_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("dynamic_update_slice", err))?;
    let update_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(update_data),
            update_shape,
            update_strides,
            update_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("dynamic_update_slice", err))?;
    let start_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            index_dtype,
            typed_bytes(starts.values.as_slice()),
            &starts.shape,
            &start_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("dynamic_update_slice", err))?;
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            operand_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("dynamic_update_slice", err))?;
    let operand_ptr = ErasedRawStridedPtr::from_ref(&operand_ref);
    let update_ptr = ErasedRawStridedPtr::from_ref(&update_ref);
    let start_ptr = ErasedRawStridedPtr::from_ref(&start_ref);
    plan.execute_uninit(&mut out_ref, &operand_ptr, &update_ptr, &start_ptr, exec)
        .map_err(|err| crate::Error::backend_source("dynamic_update_slice", err))?;

    // SAFETY: the dynamic-update plan first copies the initialized operand and
    // then overwrites the update window before returning success.
    unsafe { out.assume_init_recycled() }
}

fn typed_pad<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    input: &TypedTensor<T>,
    config: &PadConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    typed_pad_at(buffers, Operand::from_typed("pad", input)?, config, exec)
}

fn typed_pad_view<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    config: &PadConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + TensorScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_pad_at(buffers, operand_from_view(view)?, config, exec)
}

fn typed_pad_at<T: Copy + Clone + PoolScalar + TensorScalar>(
    buffers: &mut BufferPool,
    operand: Operand<'_, T>,
    config: &PadConfig,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>> {
    let Operand {
        data,
        shape: input_shape,
        strides: input_strides,
        offset: input_offset,
    } = operand;
    let rank = input_shape.len();
    if config.edge_padding_low.len() != rank {
        return Err(crate::Error::rank_mismatch(
            "pad",
            rank,
            config.edge_padding_low.len(),
        ));
    }
    if config.edge_padding_high.len() != rank {
        return Err(crate::Error::rank_mismatch(
            "pad",
            rank,
            config.edge_padding_high.len(),
        ));
    }
    if config.interior_padding.len() != rank {
        return Err(crate::Error::rank_mismatch(
            "pad",
            rank,
            config.interior_padding.len(),
        ));
    }

    let mut out_shape = Vec::with_capacity(input_shape.len());
    for (axis, &input_extent) in input_shape.iter().enumerate() {
        if config.interior_padding[axis] < 0 {
            return Err(crate::Error::invalid_argument(
                "pad",
                "configuration",
                format!("interior padding must be non-negative on axis {axis}"),
            ));
        }
        let input_extent_i64 = i64::try_from(input_extent).map_err(|_| {
            crate::Error::invalid_argument(
                "pad",
                "configuration",
                format!("input extent on axis {axis} does not fit in i64"),
            )
        })?;
        let spacing = config.interior_padding[axis]
            .checked_add(1)
            .ok_or_else(|| {
                crate::Error::invalid_argument(
                    "pad",
                    "configuration",
                    format!("interior padding overflow on axis {axis}"),
                )
            })?;
        let base = if input_extent == 0 {
            0
        } else {
            input_extent_i64
                .checked_sub(1)
                .and_then(|extent| extent.checked_mul(spacing))
                .and_then(|extent| extent.checked_add(1))
                .ok_or_else(|| {
                    crate::Error::invalid_argument(
                        "pad",
                        "configuration",
                        format!("padded input extent overflow on axis {axis}"),
                    )
                })?
        };
        let dim = config.edge_padding_low[axis]
            .checked_add(config.edge_padding_high[axis])
            .and_then(|edge| edge.checked_add(base))
            .ok_or_else(|| {
                crate::Error::invalid_argument(
                    "pad",
                    "configuration",
                    format!("output dimension overflow on axis {axis}"),
                )
            })?;
        out_shape.push(usize::try_from(dim).map_err(|_| {
            crate::Error::invalid_argument(
                "pad",
                "configuration",
                format!("negative output dimension on axis {axis}"),
            )
        })?);
    }

    let dtype = kernel_dtype(T::dtype());
    let out_strides = inline_col_major_strides("pad", &out_shape)?;
    let plan = ErasedPadPlan::compile(
        dtype,
        input_shape,
        input_strides,
        &out_shape,
        &out_strides,
        &config.edge_padding_low,
        &config.edge_padding_high,
        &config.interior_padding,
    )
    .map_err(|err| crate::Error::backend_source("pad", err))?;
    let fill = T::pool_zero();
    let mut out = PooledUninitOutput::new(buffers, &out_shape)?;
    let input_ref = unsafe {
        // SAFETY: this operation supplies initialized typed bytes with
        // alignment and dtype matching the descriptor; validated dimensions,
        // strides, and offset keep every reachable source element in bounds
        // for the retained borrow.
        crate::erased_raw_strided_ref(
            dtype,
            typed_bytes(data),
            input_shape,
            input_strides,
            input_offset,
        )
    }
    .map_err(|err| crate::Error::backend_source("pad", err))?;
    let input_ptr = ErasedRawStridedPtr::from_ref(&input_ref);
    let mut out_ref = unsafe {
        // SAFETY: the output guard exclusively owns aligned storage whose
        // byte layout agrees with dtype; validated dimensions, strides, and
        // offset keep every destination in bounds, and the following kernel
        // overwrites every reachable element before typed exposure.
        crate::erased_raw_strided_uninit_mut(
            dtype,
            out.as_uninit_bytes_mut(),
            &out_shape,
            &out_strides,
            0,
        )
    }
    .map_err(|err| crate::Error::backend_source("pad", err))?;
    plan.execute_uninit(
        &mut out_ref,
        &input_ptr,
        typed_bytes(std::slice::from_ref(&fill)),
        exec,
    )
    .map_err(|err| crate::Error::backend_source("pad", err))?;

    // INVARIANT: the compact output has no unreachable storage, so successful
    // pad replay initializes every element.
    // SAFETY: every element contains a valid T after successful replay.
    unsafe { out.assume_init_recycled() }
}

#[cfg(test)]
mod tests;