use num_complex::{Complex32, Complex64};
use tenferro_tensor_core::{promote_in_set, DType, DynRank};
use crate::{Host, ScalarSet, StridedSliceSpec, TensorScalar, TypedTensor, TypedTensorView};
fn dtype_mismatch(op: &'static str, expected: DType, actual: DType) -> crate::Error {
crate::Error::validation(
op,
tenferro_tensor_core::ValidationError::DTypeMismatch { expected, actual },
)
}
#[derive(Clone, Debug)]
pub struct DefaultScalars {
value: DefaultScalarsValue,
}
#[derive(Clone, Debug)]
pub(crate) enum DefaultScalarsValue {
F32(TypedTensor<f32, DynRank, Host>),
F64(TypedTensor<f64, DynRank, Host>),
I32(TypedTensor<i32, DynRank, Host>),
I64(TypedTensor<i64, DynRank, Host>),
Bool(TypedTensor<bool, DynRank, Host>),
C32(TypedTensor<Complex32, DynRank, Host>),
C64(TypedTensor<Complex64, DynRank, Host>),
}
impl ScalarSet for DefaultScalarsValue {
type Tag = DType;
const TAGS: &'static [Self::Tag] = DType::TAGS;
fn tag(&self) -> Self::Tag {
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,
}
}
fn promote(lhs: Self::Tag, rhs: Self::Tag) -> Self::Tag {
if matches!(lhs, DType::External(_)) {
return lhs;
}
if matches!(rhs, DType::External(_)) {
return rhs;
}
promote_in_set(DType::TAGS, DType::SPECS, lhs, rhs)
}
}
impl DefaultScalarsValue {
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.shape().contains(&0),
Self::F64(t) => t.shape().contains(&0),
Self::I32(t) => t.shape().contains(&0),
Self::I64(t) => t.shape().contains(&0),
Self::Bool(t) => t.shape().contains(&0),
Self::C32(t) => t.shape().contains(&0),
Self::C64(t) => t.shape().contains(&0),
}
}
pub fn as_view(&self) -> DefaultScalarsView<'_> {
match self {
Self::F32(t) => DefaultScalarsView::F32(t.as_view()),
Self::F64(t) => DefaultScalarsView::F64(t.as_view()),
Self::I32(t) => DefaultScalarsView::I32(t.as_view()),
Self::I64(t) => DefaultScalarsView::I64(t.as_view()),
Self::Bool(t) => DefaultScalarsView::Bool(t.as_view()),
Self::C32(t) => DefaultScalarsView::C32(t.as_view()),
Self::C64(t) => DefaultScalarsView::C64(t.as_view()),
}
}
}
impl DefaultScalars {
pub fn from_vec_col_major<T: TensorScalar>(
shape: impl Into<Vec<usize>>,
data: Vec<T>,
) -> crate::Result<Self> {
T::into_default_scalars(shape.into(), data)
}
pub fn dtype(&self) -> DType {
self.value.dtype()
}
pub fn shape(&self) -> &[usize] {
self.value.shape()
}
pub fn rank(&self) -> usize {
self.value.rank()
}
pub fn is_empty(&self) -> bool {
self.value.is_empty()
}
pub fn as_slice<T: TensorScalar>(&self) -> crate::Result<&[T]> {
let actual = self.dtype();
T::default_scalars_slice(self)
.ok_or_else(|| dtype_mismatch("DefaultScalars::as_slice", T::dtype(), actual))
}
pub fn as_mut_slice<T: TensorScalar>(&mut self) -> crate::Result<&mut [T]> {
let actual = self.dtype();
T::default_scalars_slice_mut(self)
.ok_or_else(|| dtype_mismatch("DefaultScalars::as_mut_slice", T::dtype(), actual))
}
pub fn as_view(&self) -> DefaultScalarsView<'_> {
self.value.as_view()
}
pub fn into_vec_col_major<T: TensorScalar>(self) -> crate::Result<(Vec<usize>, Vec<T>)> {
let actual = self.dtype();
T::from_default_scalars(self)
.map(TypedTensor::<T, DynRank, Host>::into_vec_col_major)
.ok_or_else(|| dtype_mismatch("DefaultScalars::into_vec_col_major", T::dtype(), actual))
}
pub(crate) fn from_payload(value: DefaultScalarsValue) -> Self {
Self { value }
}
pub(crate) fn payload(&self) -> &DefaultScalarsValue {
&self.value
}
pub(crate) fn payload_mut(&mut self) -> &mut DefaultScalarsValue {
&mut self.value
}
pub(crate) fn into_payload(self) -> DefaultScalarsValue {
self.value
}
}
impl ScalarSet for DefaultScalars {
type Tag = DType;
const TAGS: &'static [Self::Tag] = <DefaultScalarsValue as ScalarSet>::TAGS;
fn tag(&self) -> Self::Tag {
self.value.tag()
}
fn promote(lhs: Self::Tag, rhs: Self::Tag) -> Self::Tag {
<DefaultScalarsValue as ScalarSet>::promote(lhs, rhs)
}
}
#[derive(Clone, Debug)]
pub enum DefaultScalarsView<'a> {
F32(TypedTensorView<'a, f32, DynRank, Host>),
F64(TypedTensorView<'a, f64, DynRank, Host>),
I32(TypedTensorView<'a, i32, DynRank, Host>),
I64(TypedTensorView<'a, i64, DynRank, Host>),
Bool(TypedTensorView<'a, bool, DynRank, Host>),
C32(TypedTensorView<'a, Complex32, DynRank, Host>),
C64(TypedTensorView<'a, Complex64, DynRank, Host>),
}
macro_rules! impl_dynamic_view {
($self:ident, $method:ident($($arg:ident),*) => $inner:ident) => {
match $self {
DefaultScalarsView::F32(view) => DefaultScalarsView::F32(view.$method($($arg),*)?),
DefaultScalarsView::F64(view) => DefaultScalarsView::F64(view.$method($($arg),*)?),
DefaultScalarsView::I32(view) => DefaultScalarsView::I32(view.$method($($arg),*)?),
DefaultScalarsView::I64(view) => DefaultScalarsView::I64(view.$method($($arg),*)?),
DefaultScalarsView::Bool(view) => DefaultScalarsView::Bool(view.$method($($arg),*)?),
DefaultScalarsView::C32(view) => DefaultScalarsView::C32(view.$method($($arg),*)?),
DefaultScalarsView::C64(view) => DefaultScalarsView::C64(view.$method($($arg),*)?),
}
};
}
impl<'a> DefaultScalarsView<'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.shape().contains(&0),
Self::F64(view) => view.shape().contains(&0),
Self::I32(view) => view.shape().contains(&0),
Self::I64(view) => view.shape().contains(&0),
Self::Bool(view) => view.shape().contains(&0),
Self::C32(view) => view.shape().contains(&0),
Self::C64(view) => view.shape().contains(&0),
}
}
pub fn reshape_view(&self, shape: impl AsRef<[usize]>) -> crate::Result<Self> {
let shape = shape.as_ref();
Ok(impl_dynamic_view!(self, reshape_view(shape) => view))
}
pub fn transpose_view(&self, axes: &[usize]) -> crate::Result<Self> {
Ok(impl_dynamic_view!(self, transpose_view(axes) => view))
}
pub fn slice_view(&self, spec: &[StridedSliceSpec]) -> crate::Result<Self> {
Ok(impl_dynamic_view!(self, slice_view(spec) => view))
}
}
#[allow(clippy::large_enum_variant)]
#[derive(Clone, Debug)]
pub enum DefaultScalarsRef<'a> {
Tensor(&'a DefaultScalars),
View(DefaultScalarsView<'a>),
}
impl<'a> DefaultScalarsRef<'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.shape().contains(&0),
}
}
}