tokitai-operator 0.1.0

//! The `Tensor<T>` data structure.
//!
//! `Tensor<T>` is the actual data: an `ObjectMeta` plus a
//! `Vec<T>`. The `Tensor::dense_cpu` constructor builds a dense
//! CPU-stored tensor; the `Tensor::padic` constructor builds a
//! p-adic scalar tensor. The i64 backend is the default; other
//! element types (f32, p-adic digits) are constructed on demand.
//!
use crate::domain::DomainId;

use super::{Dim, ObjectMeta, Representation, Shape};

#[derive(Debug, Clone, PartialEq)]
pub struct Tensor<T> {
    pub meta: ObjectMeta,
    pub data: Vec<T>,
}

impl<T> Tensor<T> {
    /// Build a dense-CPU tensor from a domain, shape, and data vector.
    /// This is the existing unchecked constructor: it does not validate
    /// that `data.len()` matches the product of static shape dimensions.
    /// Use [`Tensor::try_from_vec`] for a fallible variant.
    pub fn dense_cpu(domain: DomainId, shape: Shape, data: Vec<T>) -> Self {
        Self {
            meta: ObjectMeta::tensor(domain, shape, Representation::dense_cpu()),
            data,
        }
    }

    /// Build a dense-CPU tensor, validating that the data length matches
    /// the product of all `Dim::Static` dimensions. Non-static dimensions
    /// (Symbolic, Bounded, DataDependent) cannot be validated, so a shape
    /// containing any of those returns an `Error::Shape`.
    pub fn try_from_vec(
        domain: DomainId,
        shape: Shape,
        data: Vec<T>,
    ) -> Result<Self, crate::Error> {
        let mut expected: usize = 1;
        for dim in shape.dims.iter() {
            expected = expected
                .checked_mul(dim.value().ok_or_else(|| {
                    crate::Error::shape(format!(
                        "try_from_vec requires all-static shape, got dim {dim:?}"
                    ))
                })?)
                .ok_or_else(|| {
                    crate::Error::shape(format!(
                        "shape product overflows usize for shape {shape:?}"
                    ))
                })?;
        }
        let expected = expected;
        if data.len() != expected {
            return Err(crate::Error::shape(format!(
                "data length {} does not match static shape product {expected} for shape {shape:?}",
                data.len()
            )));
        }
        Ok(Self::dense_cpu(domain, shape, data))
    }

    /// Number of elements in the dense data vector.
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Whether the dense data vector is empty.
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Rank of the tensor (number of dimensions).
    pub fn rank(&self) -> usize {
        self.meta.shape.rank()
    }

    /// Borrow the dense data as a slice.
    pub fn as_slice(&self) -> &[T] {
        &self.data
    }
}

impl<T: Clone> Tensor<T> {
    /// Build a rank-1 vector tensor from a domain, a single dimension, and a
    /// data vector. The data vector is moved into the tensor.
    pub fn vector(domain: DomainId, dim: usize, data: Vec<T>) -> Self {
        Self::dense_cpu(domain, Shape::from(vec![dim]), data)
    }

    /// Build a rank-2 matrix tensor from a domain, two dimensions, and a
    /// row-major data vector of length `rows * cols`. Returns an error if the
    /// data length does not match.
    pub fn matrix(
        domain: DomainId,
        rows: usize,
        cols: usize,
        data: Vec<T>,
    ) -> Result<Self, crate::Error> {
        let expected = rows.checked_mul(cols).ok_or_else(|| {
            crate::Error::shape(format!(
                "matrix dimensions overflow: rows={rows} cols={cols}"
            ))
        })?;
        if data.len() != expected {
            return Err(crate::Error::shape(format!(
                "matrix data length {} does not match rows*cols={expected}",
                data.len()
            )));
        }
        Ok(Self::dense_cpu(domain, Shape::from(vec![rows, cols]), data))
    }

    /// Borrow the rank-2 matrix rows as a slice of row slices. Returns None if
    /// the tensor is not rank-2.
    pub fn as_matrix_rows(&self) -> Option<&[T]> {
        if self.rank() == 2 {
            Some(&self.data)
        } else {
            None
        }
    }
}

impl<T: std::fmt::Display> std::fmt::Display for Tensor<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // For scalar tensors, just print the single element.
        if self.meta.shape.dims.is_empty() {
            if let Some(value) = self.data.first() {
                return write!(f, "{value}");
            }
            return write!(f, "<empty scalar>");
        }
        // For 1-D tensors, print the values comma-separated in brackets.
        if self.meta.shape.dims.len() == 1 {
            let parts: Vec<String> = self.data.iter().map(|v| v.to_string()).collect();
            return write!(f, "[{}]", parts.join(", "));
        }
        // For rank >= 2 tensors, print a small summary header.
        write!(
            f,
            "Tensor<{}>(shape={}, len={})",
            std::any::type_name::<T>(),
            self.meta.shape,
            self.data.len()
        )
    }
}

/// Iterate over owned tensor data.
impl<T> IntoIterator for Tensor<T> {
    type Item = T;
    type IntoIter = std::vec::IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        self.data.into_iter()
    }
}

/// Build a rank-1 dense i64 tensor from any iterator of i64 values.
impl FromIterator<i64> for Tensor<i64> {
    fn from_iter<I: IntoIterator<Item = i64>>(iter: I) -> Self {
        let data: Vec<i64> = iter.into_iter().collect();
        let dim = data.len();
        Self::vector(crate::domain::DomainId::new("integer"), dim, data)
    }
}

impl<T: Clone> Tensor<T> {
    /// Build a dense CPU tensor filled with the supplied constant value, with
    /// the supplied shape. Total element count is the product of the static
    /// dimensions of the shape.
    pub fn filled(domain: DomainId, shape: Shape, value: T) -> Self {
        // For a scalar (empty dims) the total is 1; otherwise the product
        // of the Static dimensions. Non-Static dimensions are skipped, so a
        // shape that mixes Static and Symbolic dims uses the Static product.
        let total: usize = if shape.dims.is_empty() {
            1
        } else {
            shape
                .dims
                .iter()
                .filter_map(|dim| match dim {
                    super::Dim::Static(value) => Some(*value),
                    _ => None,
                })
                .product()
        };
        let data: Vec<T> = (0..total).map(|_| value.clone()).collect();
        Self::dense_cpu(domain, shape, data)
    }
}

impl Tensor<i64> {
    /// Build a dense CPU i64 tensor of the supplied shape filled with zeros.
    /// Requires that all dimensions are Static; returns an error otherwise.
    pub fn zeros(domain: DomainId, shape: Shape) -> Result<Self, crate::Error> {
        for dim in &shape.dims {
            if !matches!(dim, super::Dim::Static(_)) {
                return Err(crate::Error::shape(format!(
                    "Tensor::zeros requires all-Static dimensions, got {shape:?}"
                )));
            }
        }
        let total: usize = if shape.dims.is_empty() {
            1
        } else {
            shape
                .dims
                .iter()
                .filter_map(|dim| match dim {
                    super::Dim::Static(value) => Some(*value),
                    _ => None,
                })
                .product()
        };
        Ok(Self::dense_cpu(domain, shape, vec![0_i64; total]))
    }
}

impl Tensor<i64> {
    /// Concatenate two rank-1 dense i64 tensors along their first (and only)
    /// axis. Both inputs must be rank-1 dense i64 tensors on the same domain.
    /// Returns a new rank-1 tensor with `a.len() + b.len()` elements. The
    /// resulting shape is `[a.len() + b.len()]`.
    pub fn concat_along_first_axis(a: &Tensor<i64>, b: &Tensor<i64>) -> Result<Self, crate::Error> {
        if a.rank() != 1 || b.rank() != 1 {
            return Err(crate::Error::shape(format!(
                "concat_along_first_axis requires rank-1 inputs, got ranks {} and {}",
                a.rank(),
                b.rank()
            )));
        }
        if a.meta.domain != b.meta.domain {
            return Err(crate::Error::domain(format!(
                "concat_along_first_axis requires matching domains, got {:?} and {:?}",
                a.meta.domain, b.meta.domain
            )));
        }
        let mut data: Vec<i64> = Vec::with_capacity(a.data.len() + b.data.len());
        data.extend_from_slice(&a.data);
        data.extend_from_slice(&b.data);
        let dim = data.len();
        Ok(Self::vector(a.meta.domain.clone(), dim, data))
    }
}

impl<T: PartialEq> Tensor<T> {
    /// Rank-aware element equality at a multi-dimensional index.
    ///
    /// Returns `true` when:
    /// - the linear `index` is in-bounds for `self` and `other`,
    /// - the flattened rank of `self` and `other` match,
    /// - the per-dimension sizes of `self` and `other` match, and
    /// - the element at `index` in `self` equals the element at `index` in
    ///   `other`.
    ///
    /// Returns `false` for any out-of-bounds index, shape mismatch, or value
    /// mismatch. The helper never panics, even when the index is past the
    /// end of either tensor.
    pub fn element_equals(&self, other: &Tensor<T>, index: usize) -> bool {
        // Shape must match element-wise. We compare the flat dimension list
        // because we do not want this helper to require the two tensors to
        // share the same ObjectMeta.
        if self.meta.shape.dims != other.meta.shape.dims {
            return false;
        }
        // Rank must also be the same as the data length when both are dense.
        if self.data.is_empty() && other.data.is_empty() {
            return index == 0 && self.rank() == other.rank();
        }
        if self.data.is_empty() || other.data.is_empty() {
            return false;
        }
        let self_len = self.data.len();
        let other_len = other.data.len();
        if index >= self_len || index >= other_len {
            return false;
        }
        self.data[index] == other.data[index]
    }
}

impl<T> Tensor<T> {
    /// Borrow the first element of the dense data, or `None` when empty.
    pub fn first(&self) -> Option<&T> {
        self.data.first()
    }

    /// Borrow the last element of the dense data, or `None` when empty.
    pub fn last(&self) -> Option<&T> {
        self.data.last()
    }
}

impl Tensor<i64> {
    /// Flatten a rank-2 matrix tensor into a rank-1 vector tensor by
    /// concatenating the rows in row-major order. Returns an error if the
    /// tensor is not rank-2.
    pub fn flatten_first_axis(t: &Tensor<i64>) -> Result<Tensor<i64>, crate::Error> {
        if t.rank() != 2 {
            return Err(crate::Error::verification(format!(
                "flatten_first_axis requires rank-2 tensor, got rank={}",
                t.rank()
            )));
        }
        let rows = t.meta.shape.first_dim().and_then(Dim::value).unwrap_or(0);
        let new_shape = if rows == 0 {
            Shape::from(vec![0usize])
        } else {
            // Total element count comes from data length when shape is dense.
            Shape::from(vec![t.data.len()])
        };
        Ok(Tensor {
            meta: ObjectMeta::tensor(
                t.meta.domain.clone(),
                new_shape,
                t.meta.representation.clone(),
            ),
            data: t.data.clone(),
        })
    }
}

impl Tensor<i64> {
    /// Sum of all elements in the dense data, computed with `i64::wrapping_add`
    /// so that overflow wraps rather than panicking. An empty tensor returns
    /// `0`. The sum is independent of the tensor's shape.
    pub fn sum(&self) -> i64 {
        self.data.iter().copied().fold(0i64, i64::wrapping_add)
    }

    /// Remove the first `n` elements from the dense data and return them
    /// as a `Vec<T>`. The tensor is shortened by `n` (Vec::drain semantics).
    /// Use case: popping MoE padding rows after a forward pass.
    ///
    /// Silently saturates: if `n > self.data.len()`, only `self.data.len()`
    /// items are drained (no panic). The `T: Clone` bound exists because
    /// `Vec::drain` requires the items to be movable.
    pub fn truncate(&mut self, n: usize) -> Vec<i64> {
        let drain_n = n.min(self.data.len());
        self.data.drain(..drain_n).collect()
    }

    /// Consume the tensor and return its `(meta, data)` parts. Useful
    /// when a downstream consumer needs the raw shape/meta/data triple
    /// without going through the Tensor wrapper.
    pub fn into_inner(self) -> (crate::object::ObjectMeta, Vec<i64>) {
        (self.meta, self.data)
    }

    /// Reverse the order of elements in the dense data, returning a new
    /// tensor with the same shape and domain but reversed element order.
    /// An empty tensor returns a copy with empty data.
    pub fn reversed(&self) -> Tensor<i64> {
        let mut data = self.data.clone();
        data.reverse();
        Tensor {
            meta: self.meta.clone(),
            data,
        }
    }
}

impl Tensor<i64> {
    /// Smallest element in the dense data, or `None` when empty.
    pub fn min(&self) -> Option<i64> {
        self.data.iter().copied().min()
    }

    /// Largest element in the dense data, or `None` when empty.
    pub fn max(&self) -> Option<i64> {
        self.data.iter().copied().max()
    }
}