tenferro-cpu 0.8.0

use num_complex::{Complex32, Complex64};
use num_traits::Zero;
use std::mem::MaybeUninit;
use strided_kernel::{
    col_major_strides, copy_into, copy_into_uninit, map_into, Identity, StridedView, StridedViewMut,
};

use crate::buffer_pool::{BufferPool, PoolScalar};
use tenferro_tensor::{
    DType, MemoryKind, Placement, Tensor, TensorRank, TensorRead, TensorScalar, TensorView,
    TypedTensor, TypedTensorView, TypedTensorViewMut,
};

/// The Rust scalar type behind a preset variant name a macro received.
macro_rules! preset_scalar {
    (F32) => {
        f32
    };
    (F64) => {
        f64
    };
    (I32) => {
        i32
    };
    (I64) => {
        i64
    };
    (Bool) => {
        bool
    };
    (C32) => {
        num_complex::Complex32
    };
    (C64) => {
        num_complex::Complex64
    };
}
#[cfg(test)]
use super::tensor_from_array;
#[cfg(test)]
use super::typed_array_uninit;
use super::{
    cpu_backend_buffer_error, typed_host_data, typed_view, typed_view_from_view, PooledUninitOutput,
};
use threadteam::Exec;

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

fn validate_rank(op: &'static str, expected: usize, actual: usize) -> crate::Result<()> {
    if expected != actual {
        return Err(crate::Error::rank_mismatch(op, expected, actual));
    }
    Ok(())
}

fn validate_axis(op: &'static str, axis: usize, rank: usize) -> crate::Result<()> {
    if axis >= rank {
        return Err(crate::Error::axis_out_of_bounds(op, axis, rank));
    }
    Ok(())
}

fn validate_axes_distinct(op: &'static str, axis_a: usize, axis_b: usize) -> crate::Result<()> {
    if axis_a == axis_b {
        return Err(crate::Error::duplicate_axis(op, axis_a, "axes"));
    }
    Ok(())
}

fn checked_shape_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 validate_permutation(op: &'static str, perm: &[usize], rank: usize) -> crate::Result<()> {
    validate_rank(op, rank, perm.len())?;
    // Common ranks track seen axes in a bitmask; only very high ranks allocate.
    // A materializing transpose is a metadata-only validation path at inline
    // rank, so it must not take a heap scratch here (#2178 item 7).
    if rank <= u128::BITS as usize {
        let mut seen = 0u128;
        for &axis in perm {
            validate_axis(op, axis, rank)?;
            let bit = 1u128 << axis;
            if seen & bit != 0 {
                return Err(crate::Error::duplicate_axis(op, axis, "perm"));
            }
            seen |= bit;
        }
        return Ok(());
    }
    let mut seen = vec![false; rank];
    for &axis in perm {
        validate_axis(op, axis, rank)?;
        if seen[axis] {
            return Err(crate::Error::duplicate_axis(op, axis, "perm"));
        }
        seen[axis] = true;
    }
    Ok(())
}

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

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

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

macro_rules! dispatch_tensor_view_with_bool_special_result {
    ($view:expr, |$v:ident| $body:expr, bool |$bool_v:ident| $bool_body:expr) => {
        match $view {
            TensorView::Bool($bool_v) => Ok(Tensor::from_typed::<bool>($bool_body?)),
            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::C32($v) => Ok(Tensor::from_typed::<Complex32>($body?)),
            TensorView::C64($v) => Ok(Tensor::from_typed::<Complex64>($body?)),
        }
    };
}

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

fn host_view<'a, T: Copy + TensorScalar>(
    op: &'static str,
    tensor: &'a TypedTensor<T>,
) -> crate::Result<StridedView<'a, T, Identity>> {
    if tensor.backend_buffer().is_some() {
        return Err(cpu_backend_buffer_error(op));
    }
    // The tensor caches its compact layout, so the view borrows both metadata
    // vectors instead of allocating a pair per call.
    StridedView::new(tensor.host_data()?, tensor.shape(), tensor.strides(), 0)
        .map_err(|err| crate::Error::backend_source(op, err))
}

/// Borrow a read view as a strided view, keeping its strides and offset.
fn read_host_view<'a, T, R>(
    op: &'static str,
    view: &'a TypedTensorView<'a, T, R>,
) -> crate::Result<StridedView<'a, T, Identity>>
where
    T: Copy + TensorScalar,
    R: tenferro_tensor::TensorRank,
{
    if view.backend_buffer().is_some() {
        return Err(cpu_backend_buffer_error(op));
    }
    StridedView::new(
        view.host_storage()?,
        view.shape(),
        view.strides(),
        view.offset(),
    )
    .map_err(|err| crate::Error::backend_source(op, err))
}

/// Diagonal view of `tensor` along `(axis_a, axis_b)` following the
/// `ExtractDiag` contract: `axis_b` is removed and the diagonal lands at
/// `axis_a` (`axis_a - 1` when `axis_a > axis_b`), matching
/// `shape_infer::extract_diag_shape` and the CUDA kernel.
///
/// `StridedView::diagonal_view` instead removes the larger of the two axes, so
/// for `axis_a > axis_b` the `axis_b` axis is moved to the trailing position
/// before extracting; the strided primitive then removes it.
fn extract_diagonal_view<'a, T: Copy + TensorScalar>(
    view: &StridedView<'a, T, Identity>,
    rank: usize,
    axis_a: usize,
    axis_b: usize,
) -> crate::Result<StridedView<'a, T, Identity>> {
    if axis_a < axis_b {
        return view
            .diagonal_view(&[(axis_a, axis_b)])
            .map_err(|err| crate::Error::backend_source("extract_diagonal", err));
    }
    let mut perm: Vec<usize> = (0..rank).filter(|&axis| axis != axis_b).collect();
    perm.push(axis_b);
    view.permute(&perm)
        .map_err(|err| crate::Error::backend_source("extract_diagonal", err))?
        .diagonal_view(&[(axis_a - 1, rank - 1)])
        .map_err(|err| crate::Error::backend_source("extract_diagonal", err))
}

#[cfg(test)]
fn copy_view_to_array<T: Copy + Clone + Send + Sync + TensorScalar + 'static>(
    op: &'static str,
    mut out: strided_kernel::StridedArray<T>,
    src: &StridedView<'_, T>,
    placement: &tenferro_tensor::Placement,
) -> crate::Result<TypedTensor<T>> {
    copy_into(&mut out.view_mut(), src, Exec::serial())
        .map_err(|err| crate::Error::backend_source(op, err))?;
    tensor_from_array_with_placement(op, out, placement)
}

#[cfg(test)]
fn tensor_from_array_with_placement<T: TensorScalar + 'static>(
    _op: &'static str,
    out: strided_kernel::StridedArray<T>,
    placement: &tenferro_tensor::Placement,
) -> crate::Result<TypedTensor<T>> {
    let shape = out.dims().to_vec();
    let mut output = TypedTensor::from_vec_col_major(shape, out.into_data())?;
    output.set_placement(placement.clone());
    Ok(output)
}

pub(crate) fn typed_materialize_view_with_pool<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    op: &'static str,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T, R>>
where
    T: Copy + Clone + PoolScalar + 'static,
    R: TensorRank,
{
    if view.backend_buffer().is_some() {
        return Err(cpu_backend_buffer_error(op));
    }
    validate_cpu_host_placement(op, "source", view.placement())?;
    let src: StridedView<'_, T, Identity> = StridedView::new(
        view.host_storage()?,
        view.shape(),
        view.strides(),
        view.offset(),
    )
    .map_err(|err| crate::Error::backend_source(op, err))?;
    let mut out = PooledUninitOutput::<T>::new(buffers, view.shape())?;
    // INVARIANT: `copy_into_uninit` is the full-overwrite entry: it selects the
    // blocked permutation copy on native floats in a sequential context and
    // the element-wise map elsewhere, so the destination is written completely
    // before typed exposure either way.
    copy_into_uninit(&mut out.as_uninit_view_mut()?, &src, exec)
        .map_err(|err| crate::Error::backend_source(op, err))?;
    // SAFETY: the successful copy replay writes every logical destination element.
    let mut out = unsafe { out.assume_init_as_recycled::<R>()? };
    out.set_placement(view.placement().clone());
    Ok(out)
}

pub(crate) fn typed_copy_view_into<T, R>(
    src: &TypedTensorView<'_, T, R>,
    dst: &mut TypedTensorViewMut<'_, T, R>,
    op: &'static str,
    exec: Exec<'_>,
) -> crate::Result<()>
where
    T: Copy + Send + Sync + 'static,
    R: TensorRank,
{
    if src.shape() != dst.shape() {
        return Err(crate::Error::shape_mismatch(
            op,
            src.shape().to_vec(),
            dst.shape().to_vec(),
        ));
    }
    if let (Some(src_buffer), Some(dst_buffer)) = (src.backend_buffer(), dst.backend_buffer()) {
        if std::ptr::eq(src_buffer, dst_buffer) {
            return Err(crate::Error::invalid_argument(
                op,
                "configuration",
                "CPU copy source and destination allocations must not alias",
            ));
        }
    }
    if src.backend_buffer().is_some() || dst.backend_buffer().is_some() {
        return Err(cpu_backend_buffer_error(op));
    }
    validate_cpu_host_placement(op, "source", src.placement())?;
    validate_cpu_host_placement(op, "destination", dst.placement())?;

    let src_view: StridedView<'_, T, Identity> = StridedView::new(
        src.host_storage()?,
        src.shape(),
        src.strides(),
        src.offset(),
    )
    .map_err(|err| crate::Error::backend_source(op, err))?;
    let dst_shape = dst.shape().to_vec();
    let dst_strides = dst.strides().to_vec();
    let dst_offset = dst.offset();
    let mut dst_view = StridedViewMut::new(
        dst.host_storage_mut()?,
        &dst_shape,
        &dst_strides,
        dst_offset,
    )
    .map_err(|err| crate::Error::backend_source(op, err))?;
    copy_into(&mut dst_view, &src_view, exec).map_err(|err| crate::Error::backend_source(op, err))
}

pub(crate) fn validate_cpu_host_placement(
    op: &'static str,
    role: &'static str,
    placement: &Placement,
) -> crate::Result<()> {
    if matches!(
        placement.memory_kind,
        MemoryKind::PinnedHost | MemoryKind::UnpinnedHost
    ) {
        return Ok(());
    }
    Err(crate::Error::runtime_state(
        op,
        format!(
            "CPU backend requires {role} host placement for {op}, got {:?}",
            placement.memory_kind
        ),
    ))
}

#[cfg(test)]
pub(crate) fn transpose(input: &Tensor, perm: &[usize]) -> crate::Result<Tensor> {
    with_test_pool(|buffers| transpose_with_pool(buffers, input, perm, Exec::serial()))
}

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

pub(crate) fn transpose_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    perm: &[usize],
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => transpose_with_pool(buffers, input, perm, exec),
        TensorRead::View(input) => dispatch_tensor_view_unary_result!(input, |view| {
            typed_transpose_view_with_pool(buffers, &view, perm, exec)
        }),
    }
}

pub(crate) fn reshape(input: &Tensor, shape: &[usize]) -> crate::Result<Tensor> {
    dispatch_tensor_unary_result!(input, |t| typed_reshape(t, shape))
}

/// Whether a flat `Copy`-class region of `elements` elements pays for a fork on
/// `budget` slots whose section costs `cost(width)`: the width `strided-basic`
/// would pick (`parallel_width`) is above one. The one predicate behind both
/// reshape gates below.
fn reshape_copy_pays(
    elements: usize,
    budget: std::num::NonZeroUsize,
    cost: impl Fn(std::num::NonZeroUsize) -> std::time::Duration,
) -> bool {
    use strided_kernel::execution::KernelClass;
    threadteam::width_for_by(KernelClass::Copy.estimate(elements), 1, budget, cost).get() > 1
}

/// Whether an owned `reshape` of `input` should run the copy on the pool.
///
/// The serial host copy is kept whenever the token's width for the copy is one:
/// `strided-basic` decides a flat copy's width from its serial time (elements
/// times the `Copy` class rate) against the token's dispatch cost, so a copy
/// the strided scheduler would keep sequential also skips the pooled engine
/// entry and does not take a pool buffer. This is the in-entry decision, made
/// on the token the engine minted; the entered execution-context seam uses it
/// directly.
pub(crate) fn reshape_uses_pool(input: &Tensor, exec: &Exec<'_>) -> crate::Result<bool> {
    if exec.is_serial() {
        return Ok(false);
    }
    let elements = checked_shape_product("reshape", "input shape", input.shape())?;
    Ok(reshape_copy_pays(elements, exec.budget(), |width| {
        exec.dispatch_cost(width)
    }))
}

/// The pre-entry form of [`reshape_uses_pool`], decided without minting a token.
///
/// A token charges one of the pool's [`threadteam::DispatchCosts`] kinds, which
/// depends on the calling thread, so the width cannot be known before the
/// token exists. This asks the same predicate with the cheapest cost any kind
/// could charge (`min_cost`), so `false` here implies `false` in
/// [`reshape_uses_pool`] and the copy can skip the engine entry. `true` only
/// means "ask the token".
pub(crate) fn reshape_may_use_pool(
    input: &Tensor,
    budget: std::num::NonZeroUsize,
    costs: &threadteam::DispatchCosts,
) -> crate::Result<bool> {
    use threadteam::DispatchKind::{InPool, Outside, Team};
    if budget.get() == 1 {
        return Ok(false);
    }
    let cheapest = [Team, InPool, Outside]
        .into_iter()
        .map(|kind| costs.min_cost(kind))
        .min()
        .unwrap_or(std::time::Duration::MAX);
    let elements = checked_shape_product("reshape", "input shape", input.shape())?;
    Ok(reshape_copy_pays(elements, budget, |_| cheapest))
}

pub(crate) fn reshape_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    shape: &[usize],
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    // A caller-owned payload has no preset view, so keep its typed refusal
    // instead of converting it to a view.
    if let TensorRead::Tensor(tensor) = &input {
        if matches!(tensor.dtype(), DType::External(_)) {
            return reshape(tensor, shape);
        }
    }
    // Owned compact inputs and strided views share one materializing copy: the
    // destination element order is the logical source order, so routing the
    // owned case through the strided entry keeps the same result while letting
    // a multi-thread context run the copy on its pool.
    dispatch_tensor_view_unary_result!(input.tensor_view(), |view| {
        typed_reshape_view_with_pool(buffers, &view, shape, exec)
    })
}

#[cfg(test)]
pub(crate) fn broadcast_in_dim(
    input: &Tensor,
    shape: &[usize],
    dims: &[usize],
) -> crate::Result<Tensor> {
    with_test_pool(|buffers| {
        broadcast_in_dim_with_pool(buffers, input, shape, dims, Exec::serial())
    })
}

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

pub(crate) fn broadcast_in_dim_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    shape: &[usize],
    dims: &[usize],
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => broadcast_in_dim_with_pool(buffers, input, shape, dims, exec),
        TensorRead::View(input) => dispatch_tensor_view_unary_result!(input, |view| {
            typed_broadcast_in_dim_view_with_pool(buffers, &view, shape, dims, exec)
        }),
    }
}

/// Convert a tensor to another dtype using checked dtype conversion.
///
/// Use `TensorStructural::cast` when an explicit lossy dtype projection is
/// intended.
///
/// # Examples
///
/// ```rust
/// use tenferro_cpu::CpuEngine;
/// use tenferro_tensor::{DType, Tensor, TensorStructural};
///
/// let mut engine = CpuEngine::serial();
/// let x = Tensor::from_vec_col_major(vec![2], vec![1.0_f32, 2.0]).unwrap();
/// let y = engine.convert(&x, DType::F64).unwrap();
/// assert_eq!(y.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
/// ```
///
/// # Errors
///
/// Returns an error when the requested conversion is outside tenferro's checked
/// dtype-promotion lattice.
#[cfg(test)]
pub(crate) fn convert(input: &Tensor, to: DType) -> crate::Result<Tensor> {
    with_test_pool(|buffers| convert_with_pool(buffers, input, to))
}

#[cfg(test)]
pub(crate) fn convert_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    to: DType,
) -> crate::Result<Tensor> {
    tenferro_tensor::validate::validate_convert_dtype("convert", input.dtype(), to)?;
    cast_with_pool(buffers, input, to, Exec::serial())
}

pub(crate) fn cast_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    to: DType,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    macro_rules! converted {
        ($variant:ident, $tensor:expr, $map:expr) => {
            Ok(Tensor::from_typed::<preset_scalar!($variant)>(
                typed_convert_with_pool(buffers, $tensor, $map, exec)?,
            ))
        };
    }

    match (input.dtype(), to) {
        (DType::F32, DType::F32) => Ok(Tensor::from_typed::<f32>(
            cast_operand::<f32>(input)?.duplicate()?,
        )),
        (DType::F32, DType::F64) => converted!(F64, cast_operand::<f32>(input)?, |x| x as f64),
        (DType::F32, DType::I32) => {
            validate_real_values_cast_to_i32(cast_operand::<f32>(input)?, |x| x as f64)?;
            converted!(I32, cast_operand::<f32>(input)?, |x| x as i32)
        }
        (DType::F32, DType::I64) => {
            validate_real_values_cast_to_i64(cast_operand::<f32>(input)?, |x| x as f64)?;
            converted!(I64, cast_operand::<f32>(input)?, |x| x as i64)
        }
        (DType::F32, DType::Bool) => converted!(Bool, cast_operand::<f32>(input)?, |x| x != 0.0),
        (DType::F32, DType::C32) => {
            converted!(C32, cast_operand::<f32>(input)?, |x| Complex32::new(x, 0.0))
        }
        (DType::F32, DType::C64) => {
            converted!(C64, cast_operand::<f32>(input)?, |x| Complex64::new(
                x as f64, 0.0
            ))
        }
        (DType::F64, DType::F32) => converted!(F32, cast_operand::<f64>(input)?, |x| x as f32),
        (DType::F64, DType::F64) => Ok(Tensor::from_typed::<f64>(
            cast_operand::<f64>(input)?.duplicate()?,
        )),
        (DType::F64, DType::I32) => {
            validate_real_values_cast_to_i32(cast_operand::<f64>(input)?, |x| x)?;
            converted!(I32, cast_operand::<f64>(input)?, |x| x as i32)
        }
        (DType::F64, DType::I64) => {
            validate_real_values_cast_to_i64(cast_operand::<f64>(input)?, |x| x)?;
            converted!(I64, cast_operand::<f64>(input)?, |x| x as i64)
        }
        (DType::F64, DType::Bool) => converted!(Bool, cast_operand::<f64>(input)?, |x| x != 0.0),
        (DType::F64, DType::C32) => {
            converted!(C32, cast_operand::<f64>(input)?, |x| Complex32::new(
                x as f32, 0.0
            ))
        }
        (DType::F64, DType::C64) => {
            converted!(C64, cast_operand::<f64>(input)?, |x| Complex64::new(x, 0.0))
        }
        (DType::I32, DType::F32) => converted!(F32, cast_operand::<i32>(input)?, |x| x as f32),
        (DType::I32, DType::F64) => converted!(F64, cast_operand::<i32>(input)?, |x| x as f64),
        (DType::I32, DType::I32) => Ok(Tensor::from_typed::<i32>(
            cast_operand::<i32>(input)?.duplicate()?,
        )),
        (DType::I32, DType::I64) => converted!(I64, cast_operand::<i32>(input)?, |x| x as i64),
        (DType::I32, DType::Bool) => converted!(Bool, cast_operand::<i32>(input)?, |x| x != 0),
        (DType::I32, DType::C32) => {
            converted!(C32, cast_operand::<i32>(input)?, |x| Complex32::new(
                x as f32, 0.0
            ))
        }
        (DType::I32, DType::C64) => {
            converted!(C64, cast_operand::<i32>(input)?, |x| Complex64::new(
                x as f64, 0.0
            ))
        }
        (DType::I64, DType::F32) => converted!(F32, cast_operand::<i64>(input)?, |x| x as f32),
        (DType::I64, DType::F64) => converted!(F64, cast_operand::<i64>(input)?, |x| x as f64),
        (DType::I64, DType::I32) => converted!(I32, cast_operand::<i64>(input)?, |x| x as i32),
        (DType::I64, DType::I64) => Ok(Tensor::from_typed::<i64>(
            cast_operand::<i64>(input)?.duplicate()?,
        )),
        (DType::I64, DType::Bool) => converted!(Bool, cast_operand::<i64>(input)?, |x| x != 0),
        (DType::I64, DType::C32) => {
            converted!(C32, cast_operand::<i64>(input)?, |x| Complex32::new(
                x as f32, 0.0
            ))
        }
        (DType::I64, DType::C64) => {
            converted!(C64, cast_operand::<i64>(input)?, |x| Complex64::new(
                x as f64, 0.0
            ))
        }
        (DType::Bool, DType::F32) => converted!(F32, cast_operand::<bool>(input)?, |x| if x {
            1.0
        } else {
            0.0
        }),
        (DType::Bool, DType::F64) => converted!(F64, cast_operand::<bool>(input)?, |x| if x {
            1.0
        } else {
            0.0
        }),
        (DType::Bool, DType::I32) => {
            converted!(I32, cast_operand::<bool>(input)?, |x| if x { 1 } else { 0 })
        }
        (DType::Bool, DType::I64) => {
            converted!(I64, cast_operand::<bool>(input)?, |x| if x { 1 } else { 0 })
        }
        (DType::Bool, DType::Bool) => Ok(Tensor::from_typed::<bool>(
            cast_operand::<bool>(input)?.duplicate()?,
        )),
        (DType::Bool, DType::C32) => {
            converted!(C32, cast_operand::<bool>(input)?, |x| Complex32::new(
                if x { 1.0 } else { 0.0 },
                0.0
            ))
        }
        (DType::Bool, DType::C64) => {
            converted!(C64, cast_operand::<bool>(input)?, |x| Complex64::new(
                if x { 1.0 } else { 0.0 },
                0.0
            ))
        }
        (DType::C32, DType::F32) => converted!(F32, cast_operand::<Complex32>(input)?, |z| z.re),
        (DType::C32, DType::F64) => {
            converted!(F64, cast_operand::<Complex32>(input)?, |z| z.re as f64)
        }
        (DType::C32, DType::I32) => {
            validate_real_values_cast_to_i32(cast_operand::<Complex32>(input)?, |z| z.re as f64)?;
            converted!(I32, cast_operand::<Complex32>(input)?, |z| z.re as i32)
        }
        (DType::C32, DType::I64) => {
            validate_real_values_cast_to_i64(cast_operand::<Complex32>(input)?, |z| z.re as f64)?;
            converted!(I64, cast_operand::<Complex32>(input)?, |z| z.re as i64)
        }
        (DType::C32, DType::Bool) => converted!(Bool, cast_operand::<Complex32>(input)?, |z| z.re
            != 0.0
            || z.im != 0.0),
        (DType::C32, DType::C32) => Ok(Tensor::from_typed::<Complex32>(
            cast_operand::<Complex32>(input)?.duplicate()?,
        )),
        (DType::C32, DType::C64) => {
            converted!(C64, cast_operand::<Complex32>(input)?, |z| Complex64::new(
                z.re as f64,
                z.im as f64
            ))
        }
        (DType::C64, DType::F32) => {
            converted!(F32, cast_operand::<Complex64>(input)?, |z| z.re as f32)
        }
        (DType::C64, DType::F64) => converted!(F64, cast_operand::<Complex64>(input)?, |z| z.re),
        (DType::C64, DType::I32) => {
            validate_real_values_cast_to_i32(cast_operand::<Complex64>(input)?, |z| z.re)?;
            converted!(I32, cast_operand::<Complex64>(input)?, |z| z.re as i32)
        }
        (DType::C64, DType::I64) => {
            validate_real_values_cast_to_i64(cast_operand::<Complex64>(input)?, |z| z.re)?;
            converted!(I64, cast_operand::<Complex64>(input)?, |z| z.re as i64)
        }
        (DType::C64, DType::Bool) => converted!(Bool, cast_operand::<Complex64>(input)?, |z| z.re
            != 0.0
            || z.im != 0.0),
        (DType::C64, DType::C32) => {
            converted!(C32, cast_operand::<Complex64>(input)?, |z| Complex32::new(
                z.re as f32,
                z.im as f32
            ))
        }
        (DType::C64, DType::C64) => Ok(Tensor::from_typed::<Complex64>(
            cast_operand::<Complex64>(input)?.duplicate()?,
        )),
        // An externally defined destination has no conversion table here, so the
        // conversion rejects it explicitly rather than guessing a representation.
        (_, DType::External(_)) => Err(crate::Error::dtype_mismatch("convert", input.dtype(), to)),
        // A caller-owned payload has no conversion table here, so the conversion
        // rejects it instead of guessing a representation.
        (DType::External(_), _) => Err(crate::Error::dtype_mismatch("convert", input.dtype(), to)),
    }
}
/// The typed tensor behind `input`, or the refusal a pair outside this table reports.
///
/// Callers reach this from a match on `input.dtype()`, so `None` means the tag and the runtime dtype
/// disagree rather than a caller mistake.
fn cast_operand<T: tenferro_tensor::TensorScalar>(
    input: &Tensor,
) -> crate::Result<&TypedTensor<T>> {
    input.as_typed::<T>().ok_or_else(|| {
        crate::Error::unsupported_dtype(
            "convert",
            input.dtype(),
            "the cast table does not cover this dtype",
        )
    })
}

fn validate_real_values_cast_to_i32<S: Copy + TensorScalar>(
    tensor: &TypedTensor<S>,
    real: impl Fn(S) -> f64,
) -> crate::Result<()> {
    for &value in typed_host_data("cast", tensor)? {
        validate_real_cast_to_i32(real(value))?;
    }
    Ok(())
}

fn validate_real_values_cast_to_i64<S: Copy + TensorScalar>(
    tensor: &TypedTensor<S>,
    real: impl Fn(S) -> f64,
) -> crate::Result<()> {
    for &value in typed_host_data("cast", tensor)? {
        validate_real_cast_to_i64(real(value))?;
    }
    Ok(())
}

fn validate_real_cast_to_i32(value: f64) -> crate::Result<()> {
    if !value.is_finite() {
        return Err(invalid_cast_value(format!(
            "real value must be finite when casting to i32, got {value}"
        )));
    }
    if value < i32::MIN as f64 || value > i32::MAX as f64 {
        return Err(invalid_cast_value(format!(
            "real value {value} is out of i32 range"
        )));
    }
    Ok(())
}

fn validate_real_cast_to_i64(value: f64) -> crate::Result<()> {
    const I64_MIN_F64: f64 = -9_223_372_036_854_775_808.0;
    const I64_MAX_EXCLUSIVE_F64: f64 = 9_223_372_036_854_775_808.0;

    if !value.is_finite() {
        return Err(invalid_cast_value(format!(
            "real value must be finite when casting to i64, got {value}"
        )));
    }
    if !(I64_MIN_F64..I64_MAX_EXCLUSIVE_F64).contains(&value) {
        return Err(invalid_cast_value(format!(
            "real value {value} is out of i64 range"
        )));
    }
    Ok(())
}

fn invalid_cast_value(message: String) -> crate::Error {
    crate::Error::invalid_argument("cast", "value", message)
}

#[cfg(test)]
pub(crate) fn extract_diagonal(
    input: &Tensor,
    axis_a: usize,
    axis_b: usize,
) -> crate::Result<Tensor> {
    with_test_pool(|buffers| {
        extract_diagonal_with_pool(buffers, input, axis_a, axis_b, Exec::serial())
    })
}

pub(crate) fn extract_diagonal_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    dispatch_tensor_unary_result!(input, |t| typed_extract_diagonal_with_pool(
        buffers, t, axis_a, axis_b, exec,
    ))
}

#[cfg(test)]
pub(crate) fn embed_diagonal(
    input: &Tensor,
    axis_a: usize,
    axis_b: usize,
) -> crate::Result<Tensor> {
    with_test_pool(|buffers| {
        embed_diagonal_with_pool(buffers, input, axis_a, axis_b, Exec::serial())
    })
}

pub(crate) fn embed_diagonal_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    dispatch_tensor_unary_with_bool_special_result!(
        input,
        |t| typed_embed_diagonal_with_pool(buffers, t, axis_a, axis_b, exec),
        bool | t | typed_embed_diagonal_impl(buffers, t, axis_a, axis_b, false, exec)
    )
}

#[cfg(test)]
pub(crate) fn tril(input: &Tensor, k: i64) -> crate::Result<Tensor> {
    with_test_pool(|buffers| tril_with_pool(buffers, input, k, Exec::serial()))
}

/// Mask the lower triangle of a read, reading a borrowed view in place.
pub(crate) fn tril_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => tril_with_pool(buffers, input, k, exec),
        TensorRead::View(view) => dispatch_tensor_view_with_bool_special_result!(
            view,
            |v| typed_tril_read(buffers, &v, k, exec),
            bool | v | typed_triangular_mask_read(buffers, &v, k, false, false, exec)
        ),
    }
}

/// Mask the upper triangle of a read, reading a borrowed view in place.
pub(crate) fn triu_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => triu_with_pool(buffers, input, k, exec),
        TensorRead::View(view) => dispatch_tensor_view_with_bool_special_result!(
            view,
            |v| typed_triu_read(buffers, &v, k, exec),
            bool | v | typed_triangular_mask_read(buffers, &v, k, true, false, exec)
        ),
    }
}

/// Embed a read's diagonal, reading a borrowed view in place.
pub(crate) fn embed_diagonal_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => embed_diagonal_with_pool(buffers, input, axis_a, axis_b, exec),
        TensorRead::View(view) => dispatch_tensor_view_with_bool_special_result!(
            view,
            |v| typed_embed_diagonal_read(buffers, &v, axis_a, axis_b, exec),
            bool | v
                | typed_embed_diagonal_read_with_fill(buffers, &v, axis_a, axis_b, false, exec)
        ),
    }
}

/// Extract a read's diagonal, reading a borrowed view in place.
pub(crate) fn extract_diagonal_read_with_pool(
    buffers: &mut BufferPool,
    input: TensorRead<'_>,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    match input {
        TensorRead::Tensor(input) => {
            extract_diagonal_with_pool(buffers, input, axis_a, axis_b, exec)
        }
        TensorRead::View(view) => {
            dispatch_tensor_view_unary_result!(view, |v| typed_extract_diagonal_read(
                buffers, &v, axis_a, axis_b, exec,
            ))
        }
    }
}

pub(crate) fn tril_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    dispatch_tensor_unary_with_bool_special_result!(
        input,
        |t| typed_tril_with_pool(buffers, t, k, exec),
        bool | t | typed_triangular_mask_with_fill_pool(buffers, t, k, false, false, exec)
    )
}

#[cfg(test)]
pub(crate) fn triu(input: &Tensor, k: i64) -> crate::Result<Tensor> {
    with_test_pool(|buffers| triu_with_pool(buffers, input, k, Exec::serial()))
}

pub(crate) fn triu_with_pool(
    buffers: &mut BufferPool,
    input: &Tensor,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<Tensor> {
    dispatch_tensor_unary_with_bool_special_result!(
        input,
        |t| typed_triu_with_pool(buffers, t, k, exec),
        bool | t | typed_triangular_mask_with_fill_pool(buffers, t, k, true, false, exec)
    )
}

#[cfg(test)]
pub(crate) fn typed_transpose<T: Copy + Clone + Send + Sync + TensorScalar + 'static>(
    tensor: &TypedTensor<T>,
    perm: &[usize],
) -> crate::Result<TypedTensor<T>> {
    validate_permutation("transpose", perm, tensor.shape().len())?;
    let src = typed_view("transpose", tensor)?;
    let permuted = src
        .permute(perm)
        .map_err(|err| crate::Error::backend_source("transpose", err))?;
    // SAFETY: copy_into overwrites every output element.
    let out = unsafe { typed_array_uninit(permuted.dims()) };
    copy_view_to_array("transpose", out, &permuted, tensor.placement())
}

pub(crate) fn typed_transpose_with_pool<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    perm: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
{
    typed_transpose_view_with_pool(buffers, &tensor.as_view(), perm, exec)
}

pub(crate) fn typed_transpose_view_with_pool<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    perm: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
    R: TensorRank,
{
    validate_permutation("transpose", perm, view.shape().len())?;
    let src = typed_view_from_view("transpose", view)?;
    let permuted = src
        .permute(perm)
        .map_err(|err| crate::Error::backend_source("transpose", err))?;
    checked_shape_product("transpose", "output shape", permuted.dims())?;
    let mut out = PooledUninitOutput::<T>::new(buffers, permuted.dims())?;
    // INVARIANT: `copy_into_uninit` fully overwrites the destination and selects
    // the blocked permutation engine for a sequential native-float copy.
    copy_into_uninit(&mut out.as_uninit_view_mut()?, &permuted, exec)
        .map_err(|err| crate::Error::backend_source("transpose", err))?;
    // SAFETY: the successful transpose copy writes every logical destination element.
    let mut out = unsafe { out.assume_init_recycled()? };
    out.set_placement(view.placement().clone());
    Ok(out)
}

/// # Errors
///
/// Returns [`crate::Error::Validation`] with `ShapeMismatch` when the input and
/// output element counts differ, or `InvalidArgument` when checked shape
/// products overflow `usize` or output storage cannot be constructed.
pub fn typed_reshape<T: Clone + TensorScalar + 'static>(
    tensor: &TypedTensor<T>,
    shape: &[usize],
) -> crate::Result<TypedTensor<T>> {
    let old_n = checked_shape_product("reshape", "input shape", tensor.shape())?;
    let new_n = checked_shape_product("reshape", "output shape", shape)?;
    if old_n != new_n {
        return Err(crate::Error::shape_mismatch(
            "reshape",
            tensor.shape().to_vec(),
            shape.to_vec(),
        ));
    }
    if tensor.backend_buffer().is_some() {
        return Err(cpu_backend_buffer_error("reshape"));
    }
    // INVARIANT: `typed_reshape` returns an independently owned tensor while
    // the borrowed input remains live; sharing its move-only root would violate
    // the single-owner contract, so this explicit host duplicate is required.
    // The session entry routes large inputs through `reshape_read_with_pool`
    // instead, which copies on the configured pool; this serial fallback serves
    // inputs whose copy the token's width gate (`reshape_uses_pool`) keeps serial
    // and any context that cannot parallelize the copy.
    let mut output = TypedTensor::from_vec_col_major(shape.to_vec(), tensor.host_data()?.to_vec())?;
    output.set_placement(tensor.placement().clone());
    Ok(output)
}

pub(crate) fn typed_reshape_view_with_pool<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    shape: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
    R: TensorRank,
{
    let old_n = checked_shape_product("reshape", "input shape", view.shape())?;
    let new_n = checked_shape_product("reshape", "output shape", shape)?;
    if old_n != new_n {
        return Err(crate::Error::shape_mismatch(
            "reshape",
            view.shape().to_vec(),
            shape.to_vec(),
        ));
    }

    let src = typed_view_from_view("reshape", view)?;
    let mut out = PooledUninitOutput::<T>::new(buffers, shape)?;
    let copy_strides = col_major_strides(view.shape());
    let mut copy_target = strided_kernel::StridedViewMut::new(
        out.as_uninit_slice_mut(),
        view.shape(),
        &copy_strides,
        0,
    )
    .map_err(|err| crate::Error::backend_source("reshape", err))?;
    // INVARIANT: `copy_into_uninit` fully overwrites the destination and selects
    // the blocked permutation engine for a sequential native-float copy.
    copy_into_uninit(&mut copy_target, &src, exec)
        .map_err(|err| crate::Error::backend_source("reshape", err))?;
    // SAFETY: the successful reshape copy writes every logical destination element.
    let mut out = unsafe { out.assume_init_recycled()? };
    out.set_placement(view.placement().clone());
    Ok(out)
}

#[cfg(test)]
pub(crate) fn typed_broadcast_in_dim<T: Copy + Clone + Send + Sync + TensorScalar + 'static>(
    tensor: &TypedTensor<T>,
    shape: &[usize],
    dims: &[usize],
) -> crate::Result<TypedTensor<T>> {
    typed_broadcast_in_dim_view_impl(&tensor.as_view(), shape, dims, |shape| unsafe {
        // SAFETY: broadcast materialization writes every output element before returning.
        Ok(typed_array_uninit(shape))
    })
}

pub(crate) fn typed_broadcast_in_dim_with_pool<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    shape: &[usize],
    dims: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
{
    typed_broadcast_in_dim_view_with_pool(buffers, &tensor.as_view(), shape, dims, exec)
}

pub(crate) fn typed_broadcast_in_dim_view_with_pool<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    shape: &[usize],
    dims: &[usize],
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
    R: TensorRank,
{
    validate_rank("broadcast_in_dim", view.shape().len(), dims.len())?;
    let mut seen = vec![false; shape.len()];
    let mut base_dims = vec![1usize; shape.len()];
    let mut base_strides = vec![0isize; shape.len()];
    for (src_axis, &dst_axis) in dims.iter().enumerate() {
        validate_axis("broadcast_in_dim", dst_axis, shape.len())?;
        if seen[dst_axis] {
            return Err(crate::Error::duplicate_axis(
                "broadcast_in_dim",
                dst_axis,
                "dims",
            ));
        }
        seen[dst_axis] = true;
        let source_dim = view.shape()[src_axis];
        let target_dim = shape[dst_axis];
        if source_dim != target_dim && source_dim != 1 {
            return Err(crate::Error::shape_mismatch(
                "broadcast_in_dim",
                view.shape().to_vec(),
                shape.to_vec(),
            ));
        }
        base_dims[dst_axis] = source_dim;
        base_strides[dst_axis] = view.strides()[src_axis];
    }
    if view.backend_buffer().is_some() {
        return Err(cpu_backend_buffer_error("broadcast_in_dim"));
    }
    let base: StridedView<'_, T, Identity> = StridedView::new(
        view.host_storage()?,
        &base_dims,
        &base_strides,
        view.offset(),
    )
    .map_err(|err| crate::Error::backend_source("broadcast_in_dim", err))?;
    let broadcast = base
        .broadcast(shape)
        .map_err(|err| crate::Error::backend_source("broadcast_in_dim", err))?;
    checked_shape_product("broadcast_in_dim", "output shape", shape)?;
    let mut out = PooledUninitOutput::<T>::new(buffers, shape)?;
    map_into(
        &mut out.as_uninit_view_mut()?,
        &broadcast,
        MaybeUninit::new,
        exec,
    )
    .map_err(|err| crate::Error::backend_source("broadcast_in_dim", err))?;
    // SAFETY: the successful broadcast copy writes every logical destination element.
    let mut out = unsafe { out.assume_init_recycled()? };
    out.set_placement(view.placement().clone());
    Ok(out)
}

#[cfg(test)]
fn typed_broadcast_in_dim_view_impl<T, R>(
    view: &TypedTensorView<'_, T, R>,
    shape: &[usize],
    dims: &[usize],
    make_out: impl FnOnce(&[usize]) -> crate::Result<strided_kernel::StridedArray<T>>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + Send + Sync + TensorScalar + 'static,
    R: TensorRank,
{
    validate_rank("broadcast_in_dim", view.shape().len(), dims.len())?;
    let mut seen = vec![false; shape.len()];
    let mut base_dims = vec![1usize; shape.len()];
    let mut base_strides = vec![0isize; shape.len()];
    for (src_axis, &dst_axis) in dims.iter().enumerate() {
        validate_axis("broadcast_in_dim", dst_axis, shape.len())?;
        if seen[dst_axis] {
            return Err(crate::Error::duplicate_axis(
                "broadcast_in_dim",
                dst_axis,
                "dims",
            ));
        }
        seen[dst_axis] = true;
        let source_dim = view.shape()[src_axis];
        let target_dim = shape[dst_axis];
        if source_dim != target_dim && source_dim != 1 {
            return Err(crate::Error::shape_mismatch(
                "broadcast_in_dim",
                view.shape().to_vec(),
                shape.to_vec(),
            ));
        }
        base_dims[dst_axis] = source_dim;
        base_strides[dst_axis] = view.strides()[src_axis];
    }
    if view.backend_buffer().is_some() {
        return Err(cpu_backend_buffer_error("broadcast_in_dim"));
    }
    let base: StridedView<'_, T, Identity> = StridedView::new(
        view.host_storage()?,
        &base_dims,
        &base_strides,
        view.offset(),
    )
    .map_err(|err| crate::Error::backend_source("broadcast_in_dim", err))?;
    let broadcast: StridedView<'_, T, Identity> = base
        .broadcast(shape)
        .map_err(|err| crate::Error::backend_source("broadcast_in_dim", err))?;
    checked_shape_product("broadcast_in_dim", "output shape", shape)?;
    let mut out = make_out(shape)?;
    copy_into(&mut out.view_mut(), &broadcast, Exec::serial())
        .map_err(|err| crate::Error::backend_source("broadcast_in_dim", err))?;
    tensor_from_array_with_placement("broadcast_in_dim", out, view.placement())
}

fn typed_convert_with_pool<S, T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<S>,
    f: impl Fn(S) -> T + Sync,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    S: Copy + Send + Sync + TensorScalar,
    T: Copy + Clone + PoolScalar,
{
    let mut out = PooledUninitOutput::<T>::new(buffers, tensor.shape())?;
    map_into(
        &mut out.as_uninit_view_mut()?,
        &typed_view("convert", tensor)?,
        |x| MaybeUninit::new(f(x)),
        exec,
    )
    .map_err(|err| crate::Error::backend_source("convert", err))?;
    // SAFETY: the successful conversion map writes every logical destination element.
    unsafe { out.assume_init_recycled() }
}

#[cfg(test)]
pub(crate) fn typed_extract_diagonal<T: Copy + Clone + Send + Sync + TensorScalar>(
    tensor: &TypedTensor<T>,
    axis_a: usize,
    axis_b: usize,
) -> crate::Result<TypedTensor<T>> {
    validate_axis("extract_diagonal", axis_a, tensor.shape().len())?;
    validate_axis("extract_diagonal", axis_b, tensor.shape().len())?;
    validate_axes_distinct("extract_diagonal", axis_a, axis_b)?;

    let src = host_view("extract_diagonal", tensor)?;
    let diag = extract_diagonal_view(&src, tensor.shape().len(), axis_a, axis_b)?;
    // SAFETY: copy_into overwrites every output element.
    let mut out = unsafe { typed_array_uninit(diag.dims()) };
    copy_into(&mut out.view_mut(), &diag, Exec::serial())
        .map_err(|err| crate::Error::backend_source("extract_diagonal", err))?;
    Ok(tensor_from_array(out))
}

pub(crate) fn typed_extract_diagonal_with_pool<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar,
{
    validate_axis("extract_diagonal", axis_a, tensor.shape().len())?;
    validate_axis("extract_diagonal", axis_b, tensor.shape().len())?;
    validate_axes_distinct("extract_diagonal", axis_a, axis_b)?;

    let src = host_view("extract_diagonal", tensor)?;
    let diag = extract_diagonal_view(&src, tensor.shape().len(), axis_a, axis_b)?;
    let mut out = PooledUninitOutput::<T>::new(buffers, diag.dims())?;
    // INVARIANT: `copy_into_uninit` fully overwrites the destination and routes a
    // native-float strided copy to the permutation engine, serial or parallel.
    copy_into_uninit(&mut out.as_uninit_view_mut()?, &diag, exec)
        .map_err(|err| crate::Error::backend_source("extract_diagonal", err))?;
    // SAFETY: the successful diagonal copy writes every logical destination element.
    unsafe { out.assume_init_recycled() }
}

#[cfg(test)]
pub(crate) fn typed_embed_diagonal<T: Copy + Zero + Clone + PoolScalar + 'static>(
    tensor: &TypedTensor<T>,
    axis_a: usize,
    axis_b: usize,
) -> crate::Result<TypedTensor<T>> {
    with_test_pool(|buffers| {
        typed_embed_diagonal_impl(buffers, tensor, axis_a, axis_b, T::zero(), Exec::serial())
    })
}

pub(crate) fn typed_embed_diagonal_with_pool<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Zero + Clone + PoolScalar + 'static,
{
    typed_embed_diagonal_impl(buffers, tensor, axis_a, axis_b, T::zero(), exec)
}

fn typed_embed_diagonal_impl<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    axis_a: usize,
    axis_b: usize,
    zero: T,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
{
    validate_axis("embed_diagonal", axis_a, tensor.shape().len())?;
    if axis_b > tensor.shape().len() {
        return Err(crate::Error::axis_out_of_bounds(
            "embed_diagonal",
            axis_b,
            tensor.shape().len(),
        ));
    }

    let src = host_view("embed_diagonal", tensor)?;
    embed_diagonal_from_view(buffers, &src, axis_a, axis_b, zero, exec)
}

/// Embed a borrowed read view on a diagonal, reading it in place.
fn typed_embed_diagonal_read<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Zero + Clone + PoolScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_embed_diagonal_read_with_fill(buffers, view, axis_a, axis_b, T::zero(), exec)
}

fn typed_embed_diagonal_read_with_fill<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    axis_a: usize,
    axis_b: usize,
    zero: T,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    validate_axis("embed_diagonal", axis_a, view.shape().len())?;
    if axis_b > view.shape().len() {
        return Err(crate::Error::axis_out_of_bounds(
            "embed_diagonal",
            axis_b,
            view.shape().len(),
        ));
    }
    let src = read_host_view("embed_diagonal", view)?;
    embed_diagonal_from_view(buffers, &src, axis_a, axis_b, zero, exec)
}

/// Mask the lower triangle of a borrowed read view, reading it in place.
fn typed_tril_read<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Zero + Clone + PoolScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_triangular_mask_read(buffers, view, k, false, T::zero(), exec)
}

/// Mask the upper triangle of a borrowed read view, reading it in place.
fn typed_triu_read<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Zero + Clone + PoolScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    typed_triangular_mask_read(buffers, view, k, true, T::zero(), exec)
}

/// Mask a triangular region of a borrowed read view, reading it in place.
fn typed_triangular_mask_read<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    k: i64,
    upper: bool,
    fill: T,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    let op = if upper { "triu" } else { "tril" };
    let src = read_host_view(op, view)?;
    let mut out = triangular_mask_from_view(buffers, op, &src, k, upper, fill, exec)?;
    out.set_placement(view.placement().clone());
    Ok(out)
}

/// Extract the diagonal of a borrowed read view, reading it in place.
fn typed_extract_diagonal_read<T, R>(
    buffers: &mut BufferPool,
    view: &TypedTensorView<'_, T, R>,
    axis_a: usize,
    axis_b: usize,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
    R: tenferro_tensor::TensorRank,
{
    validate_axis("extract_diagonal", axis_a, view.shape().len())?;
    validate_axis("extract_diagonal", axis_b, view.shape().len())?;
    validate_axes_distinct("extract_diagonal", axis_a, axis_b)?;
    let src = read_host_view("extract_diagonal", view)?;
    let diag = extract_diagonal_view(&src, view.shape().len(), axis_a, axis_b)?;
    let mut out = PooledUninitOutput::<T>::new(buffers, diag.dims())?;
    copy_into_uninit(&mut out.as_uninit_view_mut()?, &diag, exec)
        .map_err(|err| crate::Error::backend_source("extract_diagonal", err))?;
    // SAFETY: the successful diagonal copy writes every logical destination element.
    unsafe { out.assume_init_recycled() }
}

/// Embed a strided source on a diagonal in a compact pooled output.
fn embed_diagonal_from_view<T>(
    buffers: &mut BufferPool,
    src: &StridedView<'_, T, Identity>,
    axis_a: usize,
    axis_b: usize,
    zero: T,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
{
    let shape = src.dims();
    let mut out_shape: Vec<usize> = shape.to_vec();
    out_shape.insert(axis_b, shape[axis_a]);
    checked_shape_product("embed_diagonal", "output shape", &out_shape)?;
    let mut out = PooledUninitOutput::<T>::new(buffers, out_shape)?;
    strided_kernel::embed_diagonal_view_into_uninit(
        out.as_uninit_slice_mut(),
        src,
        axis_a,
        axis_b,
        zero,
        exec,
    )
    .map_err(|err| crate::Error::backend_source("embed_diagonal", err))?;
    // SAFETY: strided initialized the whole output, including off-diagonal zeros.
    unsafe { out.assume_init_recycled() }
}

#[cfg(test)]
pub(crate) fn typed_tril<T: Copy + Zero + Clone + TensorScalar>(
    tensor: &TypedTensor<T>,
    k: i64,
) -> crate::Result<TypedTensor<T>> {
    typed_triangular_mask(tensor, k, false)
}

pub(crate) fn typed_tril_with_pool<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Zero + Clone + PoolScalar + 'static,
{
    typed_triangular_mask_with_fill_pool(buffers, tensor, k, false, T::zero(), exec)
}

#[cfg(test)]
pub(crate) fn typed_triu<T: Copy + Zero + Clone + TensorScalar>(
    tensor: &TypedTensor<T>,
    k: i64,
) -> crate::Result<TypedTensor<T>> {
    typed_triangular_mask(tensor, k, true)
}

pub(crate) fn typed_triu_with_pool<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    k: i64,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Zero + Clone + PoolScalar + 'static,
{
    typed_triangular_mask_with_fill_pool(buffers, tensor, k, true, T::zero(), exec)
}

#[cfg(test)]
fn typed_triangular_mask<T: Copy + Zero + Clone + TensorScalar>(
    tensor: &TypedTensor<T>,
    k: i64,
    upper: bool,
) -> crate::Result<TypedTensor<T>> {
    let op = if upper { "triu" } else { "tril" };
    if tensor.shape().len() < 2 {
        return Err(crate::Error::rank_mismatch(op, 2, tensor.shape().len()));
    }

    let rows = tensor.shape()[0];
    let cols = tensor.shape()[1];
    if tensor.shape().contains(&0) {
        return tensor.duplicate();
    }

    let (batch_count, block_size) = checked_triangular_extent(op, tensor.shape(), rows, cols)?;
    let mut out = tensor.duplicate()?;
    let data = out.host_data_mut()?;

    // Intentionally sequential: triangular masks are index-dependent in the
    // innermost matrix plane and remain a dedicated CPU-kernel exception.
    for batch_idx in 0..batch_count {
        for col in 0..cols {
            let boundary = col as i128 - k as i128;
            for row in 0..rows {
                let row_idx = row;
                let row = row_idx as i128;
                let keep = if upper {
                    row <= boundary
                } else {
                    row >= boundary
                };
                if !keep {
                    let offset =
                        checked_triangular_offset(op, batch_idx, block_size, col, rows, row_idx)?;
                    data[offset] = T::zero();
                }
            }
        }
    }

    Ok(out)
}

fn typed_triangular_mask_with_fill_pool<T>(
    buffers: &mut BufferPool,
    tensor: &TypedTensor<T>,
    k: i64,
    upper: bool,
    fill: T,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
{
    let op = if upper { "triu" } else { "tril" };
    if tensor.shape().len() < 2 {
        return Err(crate::Error::rank_mismatch(op, 2, tensor.shape().len()));
    }
    if tensor.shape().contains(&0) {
        return tensor.duplicate();
    }

    let src = host_view(op, tensor)?;
    let mut out = triangular_mask_from_view(buffers, op, &src, k, upper, fill, exec)?;
    out.set_placement(tensor.placement().clone());
    Ok(out)
}

/// Mask a triangular region of `src` into a compact pooled output.
///
/// The source layout is `src` itself, so a borrowed view is read in place.
fn triangular_mask_from_view<T>(
    buffers: &mut BufferPool,
    op: &'static str,
    src: &StridedView<'_, T, Identity>,
    k: i64,
    upper: bool,
    fill: T,
    exec: Exec<'_>,
) -> crate::Result<TypedTensor<T>>
where
    T: Copy + Clone + PoolScalar + 'static,
{
    let shape = src.dims();
    if shape.len() < 2 {
        return Err(crate::Error::rank_mismatch(op, 2, shape.len()));
    }
    if shape.contains(&0) {
        return TypedTensor::<T>::from_vec_col_major(shape.to_vec(), Vec::new());
    }
    let mut out = PooledUninitOutput::<T>::new(buffers, shape)?;
    strided_kernel::triangular_mask_view_into_uninit(
        out.as_uninit_slice_mut(),
        src,
        k,
        upper,
        fill,
        exec,
    )
    .map_err(|err| crate::Error::backend_source(op, err))?;
    // SAFETY: the kernel writes every kept and masked output element.
    unsafe { out.assume_init_recycled() }
}

#[cfg(test)]
fn checked_triangular_extent(
    op: &'static str,
    shape: &[usize],
    rows: usize,
    cols: usize,
) -> crate::Result<(usize, usize)> {
    let batch_count = shape[2..].iter().try_fold(1usize, |acc, &dim| {
        acc.checked_mul(dim).ok_or_else(|| {
            crate::Error::invalid_argument(
                op,
                "configuration",
                format!("batch extent overflows usize: {acc} * {dim}"),
            )
        })
    })?;
    let block_size = rows.checked_mul(cols).ok_or_else(|| {
        crate::Error::invalid_argument(
            op,
            "configuration",
            format!("matrix block size overflows usize: {rows} * {cols}"),
        )
    })?;
    Ok((batch_count, block_size))
}

#[cfg(test)]
fn checked_triangular_offset(
    op: &'static str,
    batch_idx: usize,
    block_size: usize,
    col: usize,
    rows: usize,
    row_idx: usize,
) -> crate::Result<usize> {
    let base = batch_idx.checked_mul(block_size).ok_or_else(|| {
        crate::Error::invalid_argument(
            op,
            "configuration",
            format!("batch offset overflows usize: {batch_idx} * {block_size}"),
        )
    })?;
    let col_offset = col.checked_mul(rows).ok_or_else(|| {
        crate::Error::invalid_argument(
            op,
            "configuration",
            format!("column offset overflows usize: {col} * {rows}"),
        )
    })?;
    base.checked_add(col_offset)
        .and_then(|offset| offset.checked_add(row_idx))
        .ok_or_else(|| {
            crate::Error::invalid_argument(
                op,
                "configuration",
                "triangular mask offset overflows usize".to_string(),
            )
        })
}

#[cfg(test)]
mod tests;