use num_complex::{Complex32, Complex64};
use smallvec::SmallVec;
mod layout;
mod rank;
pub use layout::TensorLayout;
pub use rank::{DynRank, Rank, TensorRank};
pub type ShapeVec = SmallVec<[usize; 8]>;
pub type StrideVec = SmallVec<[isize; 8]>;
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("shape product {expected} does not match data length {actual}")]
ShapeDataLengthMismatch { expected: usize, actual: usize },
#[error("rank mismatch: expected {expected}, actual {actual}")]
RankMismatch { expected: usize, actual: usize },
#[error("axis {axis} out of bounds for rank {rank}")]
AxisOutOfBounds { axis: usize, rank: usize },
#[error("duplicate axis {axis} in permutation")]
DuplicateAxis { axis: usize },
#[error("invalid permutation length: expected {expected}, actual {actual}")]
InvalidPermutationLength { expected: usize, actual: usize },
#[error("invalid slice step {step}; zero is invalid and this API may require a positive step")]
InvalidSliceStep { step: isize },
#[error(
"slice bounds are invalid or unsupported: start={start}, end={end}, axis_len={axis_len}"
)]
InvalidSliceBounds {
start: isize,
end: isize,
axis_len: usize,
},
#[error("reshape element-count mismatch: from {from} to {to}")]
ReshapeElementCountMismatch { from: usize, to: usize },
#[error("view is not slice-contiguous")]
NonContiguousViewAsSlice,
#[error("dtype mismatch: expected {expected:?}, actual {actual:?}")]
DTypeMismatch { expected: DType, actual: DType },
#[error("view metadata is out of borrowed-slice bounds")]
ViewOutOfBounds,
#[error("mutable tensor layout may overlap physical elements")]
OverlappingMutableLayout,
#[error("integer overflow while validating tensor metadata")]
IntegerOverflow,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum DType {
F32,
F64,
I32,
I64,
Bool,
C32,
C64,
}
pub trait TensorScalar: Copy + Clone + Send + Sync + 'static + private::Sealed {
type Real: TensorScalar;
fn dtype() -> DType;
fn into_tensor(shape: ShapeVec, data: Vec<Self>) -> Tensor;
fn tensor_slice(tensor: &Tensor) -> Option<&[Self]>;
fn tensor_mut_slice(tensor: &mut Tensor) -> Option<&mut [Self]>;
fn into_typed(tensor: Tensor) -> Option<HostTensor<Self>>;
}
mod private {
pub trait Sealed {}
impl Sealed for f32 {}
impl Sealed for f64 {}
impl Sealed for i32 {}
impl Sealed for i64 {}
impl Sealed for bool {}
impl Sealed for num_complex::Complex32 {}
impl Sealed for num_complex::Complex64 {}
}
macro_rules! impl_scalar {
($ty:ty, $real:ty, $dtype:expr, $variant:ident) => {
impl TensorScalar for $ty {
type Real = $real;
fn dtype() -> DType {
$dtype
}
fn into_tensor(shape: ShapeVec, data: Vec<Self>) -> Tensor {
Tensor::$variant(HostTensor { data, shape })
}
fn tensor_slice(tensor: &Tensor) -> Option<&[Self]> {
match tensor {
Tensor::$variant(typed) => Some(typed.as_slice()),
_ => None,
}
}
fn tensor_mut_slice(tensor: &mut Tensor) -> Option<&mut [Self]> {
match tensor {
Tensor::$variant(typed) => Some(typed.as_mut_slice()),
_ => None,
}
}
fn into_typed(tensor: Tensor) -> Option<HostTensor<Self>> {
match tensor {
Tensor::$variant(typed) => Some(typed),
_ => None,
}
}
}
};
}
impl_scalar!(f32, f32, DType::F32, F32);
impl_scalar!(f64, f64, DType::F64, F64);
impl_scalar!(i32, i32, DType::I32, I32);
impl_scalar!(i64, i64, DType::I64, I64);
impl_scalar!(bool, bool, DType::Bool, Bool);
impl_scalar!(Complex32, f32, DType::C32, C32);
impl_scalar!(Complex64, f64, DType::C64, C64);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SliceSpec {
pub start: isize,
pub end: isize,
pub step: isize,
}
#[derive(Clone, Debug, PartialEq)]
pub struct HostTensor<T> {
data: Vec<T>,
shape: ShapeVec,
}
#[derive(Clone, Debug, PartialEq)]
pub enum Tensor {
F32(HostTensor<f32>),
F64(HostTensor<f64>),
I32(HostTensor<i32>),
I64(HostTensor<i64>),
Bool(HostTensor<bool>),
C32(HostTensor<Complex32>),
C64(HostTensor<Complex64>),
}
#[derive(Clone, Debug)]
pub struct HostTensorView<'a, T> {
data: &'a [T],
shape: ShapeVec,
strides: StrideVec,
offset: isize,
}
#[derive(Clone, Debug)]
pub enum TensorView<'a> {
F32(HostTensorView<'a, f32>),
F64(HostTensorView<'a, f64>),
I32(HostTensorView<'a, i32>),
I64(HostTensorView<'a, i64>),
Bool(HostTensorView<'a, bool>),
C32(HostTensorView<'a, Complex32>),
C64(HostTensorView<'a, Complex64>),
}
#[derive(Clone, Debug)]
pub enum TensorRef<'a> {
Tensor(&'a Tensor),
View(TensorView<'a>),
}
fn checked_product(shape: &[usize]) -> Result<usize> {
shape.iter().try_fold(1usize, |acc, &dim| {
acc.checked_mul(dim).ok_or(Error::IntegerOverflow)
})
}
fn checked_shape_len(shape: &[usize], data_len: usize) -> Result<usize> {
validate_shape_metadata(shape)?;
let expected = checked_product(shape)?;
if expected != data_len {
return Err(Error::ShapeDataLengthMismatch {
expected,
actual: data_len,
});
}
Ok(expected)
}
fn validate_shape_metadata(shape: &[usize]) -> Result<()> {
checked_product(shape)?;
col_major_strides(shape)?;
Ok(())
}
fn compact_col_major_strides(shape: &[usize]) -> StrideVec {
col_major_strides(shape).expect("HostTensor shape metadata is validated at construction")
}
pub fn col_major_strides(shape: &[usize]) -> Result<StrideVec> {
let mut strides = StrideVec::new();
let mut stride = 1isize;
for &extent in shape {
strides.push(stride);
let extent = isize::try_from(extent).map_err(|_| Error::IntegerOverflow)?;
stride = stride.checked_mul(extent).ok_or(Error::IntegerOverflow)?;
}
Ok(strides)
}
fn validate_permutation(rank: usize, axes: &[usize]) -> Result<()> {
if axes.len() != rank {
return Err(Error::InvalidPermutationLength {
expected: rank,
actual: axes.len(),
});
}
let mut seen = vec![false; rank];
for &axis in axes {
if axis >= rank {
return Err(Error::AxisOutOfBounds { axis, rank });
}
if seen[axis] {
return Err(Error::DuplicateAxis { axis });
}
seen[axis] = true;
}
Ok(())
}
fn validate_view_bounds<T>(
data: &[T],
shape: &[usize],
strides: &[isize],
offset: isize,
) -> Result<()> {
layout::validate_reachable_bounds(shape, strides, offset, data.len())
}
fn is_slice_contiguous(shape: &[usize], strides: &[isize]) -> bool {
let mut expected = 1isize;
for (&extent, &stride) in shape.iter().zip(strides) {
if extent <= 1 {
continue;
}
if stride != expected {
return false;
}
let Ok(extent) = isize::try_from(extent) else {
return false;
};
let Some(next) = expected.checked_mul(extent) else {
return false;
};
expected = next;
}
true
}
impl<T> HostTensor<T> {
pub fn from_vec_col_major(shape: impl Into<ShapeVec>, data: Vec<T>) -> Result<Self> {
let shape = shape.into();
checked_shape_len(&shape, data.len())?;
Ok(Self { data, shape })
}
pub fn shape(&self) -> &[usize] {
&self.shape
}
pub fn rank(&self) -> usize {
self.shape.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn as_slice(&self) -> &[T] {
&self.data
}
pub fn as_mut_slice(&mut self) -> &mut [T] {
&mut self.data
}
pub fn as_view(&self) -> HostTensorView<'_, T> {
HostTensorView {
data: &self.data,
shape: self.shape.clone(),
strides: compact_col_major_strides(&self.shape),
offset: 0,
}
}
pub fn into_vec_col_major(self) -> (ShapeVec, Vec<T>) {
(self.shape, self.data)
}
pub fn into_reshaped(self, shape: impl Into<ShapeVec>) -> Result<Self> {
let shape = shape.into();
let from = self.data.len();
let to = checked_product(&shape)?;
if from != to {
return Err(Error::ReshapeElementCountMismatch { from, to });
}
validate_shape_metadata(&shape)?;
Ok(Self {
data: self.data,
shape,
})
}
}
impl<'a, T> HostTensorView<'a, T> {
pub fn from_slice(
shape: impl Into<ShapeVec>,
strides: impl Into<StrideVec>,
offset: isize,
data: &'a [T],
) -> Result<Self> {
let shape = shape.into();
let strides = strides.into();
validate_view_bounds(data, &shape, &strides, offset)?;
Ok(Self {
data,
shape,
strides,
offset,
})
}
pub fn shape(&self) -> &[usize] {
&self.shape
}
pub fn strides(&self) -> &[isize] {
&self.strides
}
pub fn offset(&self) -> isize {
self.offset
}
pub fn rank(&self) -> usize {
self.shape.len()
}
pub fn is_empty(&self) -> bool {
self.shape.contains(&0)
}
pub fn is_compact_col_major(&self) -> bool {
is_slice_contiguous(&self.shape, &self.strides)
}
pub fn is_zero_offset_col_major(&self) -> bool {
self.offset == 0 && self.is_compact_col_major()
}
pub fn as_slice(&self) -> Result<&'a [T]> {
if !is_slice_contiguous(&self.shape, &self.strides) {
return Err(Error::NonContiguousViewAsSlice);
}
let len = checked_product(&self.shape)?;
let start = usize::try_from(self.offset).map_err(|_| Error::IntegerOverflow)?;
let end = start.checked_add(len).ok_or(Error::IntegerOverflow)?;
self.data.get(start..end).ok_or(Error::ViewOutOfBounds)
}
pub fn reshape_view(&self, shape: impl Into<ShapeVec>) -> Result<Self> {
if !self.is_compact_col_major() {
return Err(Error::NonContiguousViewAsSlice);
}
let shape = shape.into();
let from = checked_product(&self.shape)?;
let to = checked_product(&shape)?;
if from != to {
return Err(Error::ReshapeElementCountMismatch { from, to });
}
Self::from_slice(
shape.clone(),
col_major_strides(&shape)?,
self.offset,
self.data,
)
}
pub fn transpose_view(&self, axes: &[usize]) -> Result<Self> {
validate_permutation(self.rank(), axes)?;
let shape = axes
.iter()
.map(|&axis| self.shape[axis])
.collect::<ShapeVec>();
let strides = axes
.iter()
.map(|&axis| self.strides[axis])
.collect::<StrideVec>();
Self::from_slice(shape, strides, self.offset, self.data)
}
pub fn slice_view(&self, spec: &[SliceSpec]) -> Result<Self> {
if spec.len() != self.rank() {
return Err(Error::RankMismatch {
expected: self.rank(),
actual: spec.len(),
});
}
let mut shape = ShapeVec::new();
let mut strides = StrideVec::new();
let mut offset = self.offset;
for ((&axis_len, &stride), slice) in self.shape.iter().zip(self.strides.iter()).zip(spec) {
if slice.step <= 0 {
return Err(Error::InvalidSliceStep { step: slice.step });
}
if slice.start < 0 || slice.end < 0 {
return Err(Error::InvalidSliceBounds {
start: slice.start,
end: slice.end,
axis_len,
});
}
let start = usize::try_from(slice.start).map_err(|_| Error::IntegerOverflow)?;
let end = usize::try_from(slice.end).map_err(|_| Error::IntegerOverflow)?;
if start > axis_len || end > axis_len {
return Err(Error::InvalidSliceBounds {
start: slice.start,
end: slice.end,
axis_len,
});
}
let step = usize::try_from(slice.step).map_err(|_| Error::IntegerOverflow)?;
let extent = if start >= end {
0
} else {
end.checked_sub(start)
.and_then(|span| span.checked_add(step - 1))
.ok_or(Error::IntegerOverflow)?
/ step
};
let start_offset = isize::try_from(start)
.map_err(|_| Error::IntegerOverflow)?
.checked_mul(stride)
.ok_or(Error::IntegerOverflow)?;
offset = offset
.checked_add(start_offset)
.ok_or(Error::IntegerOverflow)?;
let new_stride = stride
.checked_mul(slice.step)
.ok_or(Error::IntegerOverflow)?;
shape.push(extent);
strides.push(new_stride);
}
Self::from_slice(shape, strides, offset, self.data)
}
}
impl Tensor {
pub fn from_vec_col_major<T: TensorScalar>(
shape: impl Into<ShapeVec>,
data: Vec<T>,
) -> Result<Self> {
let shape = shape.into();
checked_shape_len(&shape, data.len())?;
Ok(T::into_tensor(shape, data))
}
pub fn dtype(&self) -> DType {
match self {
Self::F32(_) => DType::F32,
Self::F64(_) => DType::F64,
Self::I32(_) => DType::I32,
Self::I64(_) => DType::I64,
Self::Bool(_) => DType::Bool,
Self::C32(_) => DType::C32,
Self::C64(_) => DType::C64,
}
}
pub fn shape(&self) -> &[usize] {
match self {
Self::F32(t) => t.shape(),
Self::F64(t) => t.shape(),
Self::I32(t) => t.shape(),
Self::I64(t) => t.shape(),
Self::Bool(t) => t.shape(),
Self::C32(t) => t.shape(),
Self::C64(t) => t.shape(),
}
}
pub fn rank(&self) -> usize {
self.shape().len()
}
pub fn is_empty(&self) -> bool {
match self {
Self::F32(t) => t.is_empty(),
Self::F64(t) => t.is_empty(),
Self::I32(t) => t.is_empty(),
Self::I64(t) => t.is_empty(),
Self::Bool(t) => t.is_empty(),
Self::C32(t) => t.is_empty(),
Self::C64(t) => t.is_empty(),
}
}
pub fn as_slice<T: TensorScalar>(&self) -> Result<&[T]> {
T::tensor_slice(self).ok_or(Error::DTypeMismatch {
expected: T::dtype(),
actual: self.dtype(),
})
}
pub fn as_mut_slice<T: TensorScalar>(&mut self) -> Result<&mut [T]> {
let actual = self.dtype();
T::tensor_mut_slice(self).ok_or(Error::DTypeMismatch {
expected: T::dtype(),
actual,
})
}
pub fn as_view(&self) -> TensorView<'_> {
match self {
Self::F32(t) => TensorView::F32(t.as_view()),
Self::F64(t) => TensorView::F64(t.as_view()),
Self::I32(t) => TensorView::I32(t.as_view()),
Self::I64(t) => TensorView::I64(t.as_view()),
Self::Bool(t) => TensorView::Bool(t.as_view()),
Self::C32(t) => TensorView::C32(t.as_view()),
Self::C64(t) => TensorView::C64(t.as_view()),
}
}
pub fn into_vec_col_major<T: TensorScalar>(self) -> Result<(ShapeVec, Vec<T>)> {
let actual = self.dtype();
T::into_typed(self)
.map(HostTensor::into_vec_col_major)
.ok_or(Error::DTypeMismatch {
expected: T::dtype(),
actual,
})
}
}
macro_rules! impl_dynamic_view {
($self:ident, $method:ident($($arg:ident),*) => $inner:ident) => {
match $self {
TensorView::F32(view) => TensorView::F32(view.$method($($arg),*)?),
TensorView::F64(view) => TensorView::F64(view.$method($($arg),*)?),
TensorView::I32(view) => TensorView::I32(view.$method($($arg),*)?),
TensorView::I64(view) => TensorView::I64(view.$method($($arg),*)?),
TensorView::Bool(view) => TensorView::Bool(view.$method($($arg),*)?),
TensorView::C32(view) => TensorView::C32(view.$method($($arg),*)?),
TensorView::C64(view) => TensorView::C64(view.$method($($arg),*)?),
}
};
}
impl<'a> TensorView<'a> {
pub fn dtype(&self) -> DType {
match self {
Self::F32(_) => DType::F32,
Self::F64(_) => DType::F64,
Self::I32(_) => DType::I32,
Self::I64(_) => DType::I64,
Self::Bool(_) => DType::Bool,
Self::C32(_) => DType::C32,
Self::C64(_) => DType::C64,
}
}
pub fn shape(&self) -> &[usize] {
match self {
Self::F32(view) => view.shape(),
Self::F64(view) => view.shape(),
Self::I32(view) => view.shape(),
Self::I64(view) => view.shape(),
Self::Bool(view) => view.shape(),
Self::C32(view) => view.shape(),
Self::C64(view) => view.shape(),
}
}
pub fn rank(&self) -> usize {
self.shape().len()
}
pub fn is_empty(&self) -> bool {
match self {
Self::F32(view) => view.is_empty(),
Self::F64(view) => view.is_empty(),
Self::I32(view) => view.is_empty(),
Self::I64(view) => view.is_empty(),
Self::Bool(view) => view.is_empty(),
Self::C32(view) => view.is_empty(),
Self::C64(view) => view.is_empty(),
}
}
pub fn reshape_view(&self, shape: impl Into<ShapeVec>) -> Result<Self> {
let shape = shape.into();
Ok(impl_dynamic_view!(self, reshape_view(shape) => view))
}
pub fn transpose_view(&self, axes: &[usize]) -> Result<Self> {
Ok(impl_dynamic_view!(self, transpose_view(axes) => view))
}
pub fn slice_view(&self, spec: &[SliceSpec]) -> Result<Self> {
Ok(impl_dynamic_view!(self, slice_view(spec) => view))
}
}
impl<'a> TensorRef<'a> {
pub fn dtype(&self) -> DType {
match self {
Self::Tensor(tensor) => tensor.dtype(),
Self::View(view) => view.dtype(),
}
}
pub fn shape(&self) -> &[usize] {
match self {
Self::Tensor(tensor) => tensor.shape(),
Self::View(view) => view.shape(),
}
}
pub fn rank(&self) -> usize {
self.shape().len()
}
pub fn is_empty(&self) -> bool {
match self {
Self::Tensor(tensor) => tensor.is_empty(),
Self::View(view) => view.is_empty(),
}
}
}