tensor4all_tensorbackend/

tensor_element.rs

use anyhow::{anyhow, ensure, Result};
use num_complex::{Complex32, Complex64};
use tenferro::{DType, Tensor as NativeTensor, TensorScalar};

/// Public scalar element types supported by tensor4all dense/diag constructors.
///
/// Implemented for `f32`, `f64`, `Complex32`, and `Complex64`.
///
/// # Examples
///
/// ```
/// use tensor4all_tensorbackend::TensorElement;
///
/// let t = f64::dense_native_tensor_from_col_major(&[1.0, 2.0], &[2]).unwrap();
/// assert_eq!(t.shape(), &[2]);
///
/// let vals = f64::dense_values_from_native_col_major(&t).unwrap();
/// assert_eq!(vals, vec![1.0, 2.0]);
/// ```
pub trait TensorElement: TensorScalar + Copy + Send + Sync + 'static {
    /// Build a dense native tensor from column-major data.
    /// # Errors
    ///
    /// Returns an error when the shape product would overflow `usize`, when
    /// `data.len()` does not match the expected shape, or when the backend
    /// rejects the tensor.
    ///
    fn dense_native_tensor_from_col_major(data: &[Self], dims: &[usize]) -> Result<NativeTensor>;

    /// Build a dense native tensor by moving an owned column-major payload.
    ///
    /// This is the allocation-preserving counterpart of
    /// [`Self::dense_native_tensor_from_col_major`]. The payload must contain
    /// exactly the product of `dims` elements.
    ///
    /// # Errors
    ///
    /// Returns an error for a shape mismatch when the data length does not
    /// equal the dimension product, or for a backend conversion failure.
    ///
    /// # Arguments
    ///
    /// * `data` - Owned column-major values whose length equals the product of
    ///   `dims`.
    /// * `dims` - Logical tensor dimensions in column-major axis order.
    ///
    /// # Returns
    ///
    /// A native dense tensor that owns the supplied payload without cloning it.
    ///
    /// # Examples
    ///
    /// ```
    /// use tensor4all_tensorbackend::TensorElement;
    ///
    /// let tensor = f64::dense_native_tensor_from_col_major_owned(
    ///     vec![1.0, 2.0, 3.0, 4.0],
    ///     &[2, 2],
    /// )
    /// .unwrap();
    /// assert_eq!(tensor.shape(), &[2, 2]);
    /// assert_eq!(tensor.as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
    /// ```
    fn dense_native_tensor_from_col_major_owned(
        data: Vec<Self>,
        dims: &[usize],
    ) -> Result<NativeTensor>;

    /// Build a diagonal native tensor from column-major diagonal payload data.
    /// # Errors
    ///
    /// Returns an error when `logical_rank` is zero, when expanding the
    /// diagonal would overflow `usize`, or when the backend rejects the
    /// resulting dense tensor.
    ///
    fn diag_native_tensor_from_col_major(
        data: &[Self],
        logical_rank: usize,
    ) -> Result<NativeTensor>;

    /// Build a rank-0 native tensor.
    /// # Errors
    ///
    /// Returns an error when the backend rejects the rank-0 tensor as invalid
    /// input.
    ///
    fn scalar_native_tensor(value: Self) -> Result<NativeTensor>;

    /// Materialize dense column-major values from a native tensor.
    /// # Errors
    ///
    /// Returns an error when `tensor`'s dtype does not match `Self`.
    ///
    fn dense_values_from_native_col_major(tensor: &NativeTensor) -> Result<Vec<Self>>;

    /// Materialize dense column-major values while consuming the native tensor.
    ///
    /// This is the ownership-preserving counterpart of
    /// [`Self::dense_values_from_native_col_major`]: a compact column-major host
    /// result hands its buffer over instead of being copied, so a caller that
    /// built the tensor only to read it back pays no extra full-size
    /// allocation.
    ///
    /// # Arguments
    ///
    /// * `tensor` - Consumed native tensor whose logical values are returned.
    ///
    /// # Returns
    ///
    /// The tensor's values in column-major order, in `Self`'s scalar type.
    ///
    /// # Errors
    ///
    /// Returns an error on a dtype mismatch between `tensor` and `Self`, and when
    /// `tensor`'s storage cannot hand over a compact column-major host buffer
    /// because it is backend-resident or its layout is not compact, instead of
    /// copying or silently reordering it.
    ///
    /// # Examples
    ///
    /// ```
    /// use tensor4all_tensorbackend::TensorElement;
    ///
    /// let tensor =
    ///     f64::dense_native_tensor_from_col_major(&[1.0, 3.0, 2.0, 4.0], &[2, 2]).unwrap();
    /// let values = f64::into_dense_values_from_native_col_major(tensor).unwrap();
    /// assert_eq!(values, vec![1.0, 3.0, 2.0, 4.0]);
    /// ```
    fn into_dense_values_from_native_col_major(tensor: NativeTensor) -> Result<Vec<Self>>;

    /// Materialize diagonal values from a dense native tensor.
    /// # Errors
    ///
    /// Returns an error when `tensor` is rank 0, when its dimensions are not
    /// all equal, or when its dtype does not match `Self`.
    ///
    fn diag_values_from_native_temp(tensor: &NativeTensor) -> Result<Vec<Self>>;
}

fn tensor_dtype_name(dtype: DType) -> &'static str {
    match dtype {
        DType::F32 => "f32",
        DType::F64 => "f64",
        DType::I32 => "i32",
        DType::I64 => "i64",
        DType::Bool => "bool",
        DType::C32 => "c32",
        DType::C64 => "c64",
        DType::External(_) => "external",
    }
}

fn checked_product(dims: &[usize]) -> Result<usize> {
    dims.iter().try_fold(1usize, |acc, &dim| {
        acc.checked_mul(dim)
            .ok_or_else(|| anyhow::anyhow!("dimension product overflow"))
    })
}

fn dense_diagonal_values<T: Copy + Default>(diag: &[T], logical_rank: usize) -> Result<Vec<T>> {
    ensure!(
        logical_rank >= 1,
        "diagonal tensor construction requires at least one logical axis"
    );
    let diag_len = diag.len();
    let dims = vec![diag_len; logical_rank];
    let total_len = checked_product(&dims)?;
    let mut dense = vec![T::default(); total_len];
    let diagonal_stride = (0..logical_rank)
        .scan(1usize, |stride, _| {
            let current = *stride;
            *stride = stride.saturating_mul(diag_len);
            Some(current)
        })
        .sum::<usize>();
    for (i, value) in diag.iter().copied().enumerate() {
        dense[i * diagonal_stride] = value;
    }
    Ok(dense)
}

macro_rules! impl_tensor_element {
    ($ty:ty, $dtype:expr) => {
        impl TensorElement for $ty {
            fn dense_native_tensor_from_col_major(
                data: &[Self],
                dims: &[usize],
            ) -> Result<NativeTensor> {
                let expected_len: usize = checked_product(dims)?;
                ensure!(
                    data.len() == expected_len,
                    "dense tensor len {} does not match dims {:?} (expected {})",
                    data.len(),
                    dims,
                    expected_len
                );
                Ok(NativeTensor::from_vec_col_major(
                    dims.to_vec(),
                    data.to_vec(),
                )?)
            }

            fn dense_native_tensor_from_col_major_owned(
                data: Vec<Self>,
                dims: &[usize],
            ) -> Result<NativeTensor> {
                let expected_len: usize = checked_product(dims)?;
                ensure!(
                    data.len() == expected_len,
                    "dense tensor len {} does not match dims {:?} (expected {})",
                    data.len(),
                    dims,
                    expected_len
                );
                Ok(NativeTensor::from_vec_col_major(dims.to_vec(), data)?)
            }

            fn diag_native_tensor_from_col_major(
                data: &[Self],
                logical_rank: usize,
            ) -> Result<NativeTensor> {
                let dims = vec![data.len(); logical_rank];
                let dense = dense_diagonal_values(data, logical_rank)?;
                Self::dense_native_tensor_from_col_major(&dense, &dims)
            }

            fn scalar_native_tensor(value: Self) -> Result<NativeTensor> {
                Ok(NativeTensor::from_vec_col_major(vec![], vec![value])?)
            }

            fn dense_values_from_native_col_major(tensor: &NativeTensor) -> Result<Vec<Self>> {
                tensor
                    .as_slice::<Self>()
                    .map(|values| values.to_vec())
                    .map_err(|_| {
                        anyhow!(
                            "tensor dtype mismatch: expected {}, got {}",
                            tensor_dtype_name($dtype),
                            tensor_dtype_name(tensor.dtype())
                        )
                    })
            }

            fn into_dense_values_from_native_col_major(tensor: NativeTensor) -> Result<Vec<Self>> {
                let actual = tensor.dtype();
                let typed = Self::into_typed(tensor).map_err(|failure| {
                    let (_, error) = failure.into_parts();
                    anyhow!(
                        "tensor dtype mismatch: expected {}, got {}: {error}",
                        tensor_dtype_name($dtype),
                        tensor_dtype_name(actual)
                    )
                })?;
                // `into_vec_col_major` moves a compact column-major host buffer
                // and rejects a backend-resident or strided one typed, so this
                // route never copies and never silently reorders a layout.
                typed
                    .into_vec_col_major()
                    .map(|(_shape, values)| values)
                    .map_err(|failure| {
                        let (_, error) = failure.into_parts();
                        anyhow!(
                            "native {} tensor buffer cannot be transferred: {error}",
                            tensor_dtype_name($dtype)
                        )
                    })
            }

            fn diag_values_from_native_temp(tensor: &NativeTensor) -> Result<Vec<Self>> {
                let shape = tensor.shape();
                ensure!(
                    !shape.is_empty(),
                    "diagonal extraction requires rank >= 1, got scalar tensor"
                );
                let diag_len = shape[0];
                ensure!(
                    shape.iter().all(|&dim| dim == diag_len),
                    "expected square/equal dims for diagonal extraction, got {:?}",
                    shape
                );
                let dense = Self::dense_values_from_native_col_major(tensor)?;
                let diagonal_stride = (0..shape.len())
                    .scan(1usize, |stride, _| {
                        let current = *stride;
                        *stride = stride.saturating_mul(diag_len);
                        Some(current)
                    })
                    .sum::<usize>();
                Ok((0..diag_len).map(|i| dense[i * diagonal_stride]).collect())
            }
        }
    };
}

impl_tensor_element!(f32, DType::F32);
impl_tensor_element!(f64, DType::F64);
impl_tensor_element!(Complex32, DType::C32);
impl_tensor_element!(Complex64, DType::C64);

#[cfg(test)]
mod tests;