lumen_core/dynamic/
mod.rs

1mod float;
2mod integer;
3pub use float::*;
4pub use integer::*;
5use crate::{DType, Shape, Tensor, WithDType};
6use paste::paste;
7use half::bf16;
8
9#[derive(Clone)]
10pub enum DynTensor {
11    Bool(Tensor<bool>),
12    Bf16(Tensor<bf16>),
13    F32(Tensor<f32>),
14    F64(Tensor<f64>),
15    I32(Tensor<i32>),
16    U32(Tensor<u32>),
17    U8(Tensor<u8>),
18}
19
20macro_rules! method_for_all_dtype {
21    ($self:ident, $t:ident, $expr:expr) => {
22        match &$self {
23            DynTensor::Bool($t) => $expr,
24            DynTensor::Bf16($t) => $expr,
25            DynTensor::F32($t) => $expr,
26            DynTensor::F64($t) => $expr,
27            DynTensor::U32($t) => $expr,
28            DynTensor::I32($t) => $expr,
29            DynTensor::U8($t) => $expr,
30        }
31    };
32}
33
34macro_rules! method_for_all_dtype_with_dyn {
35    ($self:ident, $t:ident, $expr:expr) => {
36        match &$self {
37            DynTensor::Bool($t) => $expr.map(DynTensor::Bool),
38            DynTensor::Bf16($t) => $expr.map(DynTensor::Bf16),
39            DynTensor::F32($t) => $expr.map(DynTensor::F32),
40            DynTensor::F64($t) => $expr.map(DynTensor::F64),
41            DynTensor::U32($t) => $expr.map(DynTensor::U32),
42            DynTensor::I32($t) => $expr.map(DynTensor::I32),
43            DynTensor::U8($t) => $expr.map(DynTensor::U8),
44        }
45    };
46}
47
48impl DynTensor {
49    pub fn dtype(&self) -> DType {
50        match self {
51            Self::Bool(_) => DType::Bool,
52            Self::Bf16(_) => DType::Bf16,
53            Self::F32(_) => DType::F32,
54            Self::F64(_) => DType::F64,
55            Self::U8(_) => DType::U8,
56            Self::I32(_) => DType::I32,
57            Self::U32(_) => DType::U32,
58        }
59    }
60
61    pub fn shape(&self) -> &Shape {
62        method_for_all_dtype!(self, t, t.shape())
63    }
64
65    pub fn transpose_last(&self) -> crate::Result<Self> {
66        method_for_all_dtype_with_dyn!(self, t, t.transpose_last())
67    }
68
69    pub fn contiguous(&self) -> crate::Result<Self> {
70        method_for_all_dtype_with_dyn!(self, t, t.contiguous())
71    }
72
73    pub fn as_tensor<T: WithDType>(&self) -> crate::Result<Tensor<T>> {
74        T::from_dyn(self)
75    }
76}
77
78macro_rules! impl_convert_with_type {
79    // Rule 1: Allow exact match and specific castable variants
80    ($variant:ident, $inner:ty, [$($cast_variant:ident),*]) => {
81        paste! {
82            impl DynTensor {
83                pub fn [< is_ $inner >](&self) -> bool {
84                    matches!(self, Self::$variant(_))
85                }
86
87                pub fn [< as_ $inner >](&self) -> Option<Tensor<$inner>> {
88                    match self {
89                        Self::$variant(t) => Some(t.clone()),
90                        // Note: `as_type` usually implies exact/cheap match. 
91                        // If you want `.as_f32()` to do expensive casting, you can add it here, 
92                        // but it's usually better practice to keep `as_` exact and rely on `TryFrom`/`cast` for casting.
93                        _ => None,
94                    }
95                }
96            }
97    
98            impl From<Tensor<$inner>> for DynTensor {
99                fn from(t: Tensor<$inner>) -> Self {
100                    DynTensor::$variant(t)
101                }
102            }
103
104            impl From<&Tensor<$inner>> for DynTensor {
105                fn from(t: &Tensor<$inner>) -> Self {
106                    DynTensor::$variant(t.clone())
107                }
108            }
109    
110            impl TryFrom<DynTensor> for Tensor<$inner> {
111                type Error = crate::Error;
112    
113                #[allow(unreachable_patterns)]
114                fn try_from(value: DynTensor) -> Result<Self, Self::Error> {
115                    match value {
116                        // 1. Exact Match
117                        DynTensor::$variant(t) => Ok(t),
118                        
119                        // 2. Castable Matches
120                        // This expands to: `DynTensor::F64(t) => t.cast::<f32>(),` etc.
121                        $(
122                            DynTensor::$cast_variant(t) => t.cast::<$inner>(),
123                        )*
124                        
125                        // 3. Fallback Error
126                        // Notice we bind to `other` so we can call `.dtype()` without moving `value`
127                        other => Err(crate::Error::UnexpectedDType { 
128                            msg: "failed to cast in try_from", 
129                            expected: DType::$variant, 
130                            got: other.dtype() 
131                        }),
132                    }
133                }
134            }
135        }
136    };
137    
138    // Rule 2: Fallback for types that shouldn't accept casting from anything else (Empty array)
139    ($variant:ident, $inner:ty) => {
140        impl_convert_with_type!($variant, $inner, []);
141    };
142}
143
144impl_convert_with_type!(Bf16, bf16, [F64, F32]);
145impl_convert_with_type!(F32, f32, [F64, Bf16]);
146impl_convert_with_type!(F64, f64, [F32, Bf16]);
147
148// Types with no implicit casting allowed from DynTensor just use the standard call
149impl_convert_with_type!(Bool, bool, [U8, U32, I32, Bf16, F32, F64]);
150impl_convert_with_type!(I32, i32, [U8, U32, I32, Bf16, F32, F64]);
151impl_convert_with_type!(U32, u32, [U8, I32, I32, Bf16, F32, F64]);
152impl_convert_with_type!(U8, u8, [I32, U32, I32, Bf16, F32, F64]);
lumen_core/dynamic/mod.rs

lumen_core/dynamic/
mod.rs
