luma_tensor/dynamic/
mod.rs

1//! Type-erased tensor enum for runtime kind dispatch.
2//!
3//! [`DynTensor<D>`] erases the compile-time kind parameter (`Float`/`Int`/`Bool`)
4//! into a three-variant enum, enabling heterogeneous collections like
5//! `HashMap<String, DynTensor<D>>` for model serialization (safetensors I/O).
6
7use crate::dtype::{BoolDType, DType, KindTag};
8use crate::error::Result;
9use crate::tensor::{Shape, Tensor};
10use crate::{Bool, Device, Float, Int};
11
12/// A tensor whose kind (`Float`/`Int`/`Bool`) is determined at runtime.
13///
14/// This is the type-erased counterpart to [`Tensor<D, K>`]. Use it when you need
15/// to store or pass tensors of different kinds together — most importantly for
16/// disk I/O ([`from_bytes`](Self::from_bytes) / [`to_bytes`](Self::to_bytes)).
17pub enum DynTensor<D: Device> {
18    Float(Tensor<D, Float>),
19    Int(Tensor<D, Int>),
20    Bool(Tensor<D, Bool>),
21}
22
23// ---- accessors ----
24
25impl<D: Device> DynTensor<D> {
26    /// Runtime element type.
27    pub fn dtype(&self) -> DType {
28        match self {
29            Self::Float(t) => t.dtype().into(),
30            Self::Int(t) => t.dtype().into(),
31            Self::Bool(_) => DType::Bool,
32        }
33    }
34
35    /// Logical shape.
36    pub fn shape(&self) -> &Shape {
37        match self {
38            Self::Float(t) => t.shape(),
39            Self::Int(t) => t.shape(),
40            Self::Bool(t) => t.shape(),
41        }
42    }
43
44    /// Dimensions as a slice.
45    pub fn dims(&self) -> &[usize] {
46        self.shape().dims()
47    }
48
49    /// Device that owns the storage.
50    pub fn device(&self) -> &D {
51        match self {
52            Self::Float(t) => t.device(),
53            Self::Int(t) => t.device(),
54            Self::Bool(t) => t.device(),
55        }
56    }
57}
58
59// ---- checked conversions ----
60
61impl<D: Device> DynTensor<D> {
62    /// Borrow as a Float tensor, or `None`.
63    pub fn as_float(&self) -> Option<&Tensor<D, Float>> {
64        match self {
65            Self::Float(t) => Some(t),
66            _ => None,
67        }
68    }
69
70    /// Borrow as an Int tensor, or `None`.
71    pub fn as_int(&self) -> Option<&Tensor<D, Int>> {
72        match self {
73            Self::Int(t) => Some(t),
74            _ => None,
75        }
76    }
77
78    /// Borrow as a Bool tensor, or `None`.
79    pub fn as_bool(&self) -> Option<&Tensor<D, Bool>> {
80        match self {
81            Self::Bool(t) => Some(t),
82            _ => None,
83        }
84    }
85
86    /// Unwrap into a Float tensor, or error.
87    pub fn into_float(self) -> Result<Tensor<D, Float>> {
88        match self {
89            Self::Float(t) => Ok(t),
90            other => Err(crate::Error::Msg(format!("expected Float DynTensor, got {:?}", other.dtype()))),
91        }
92    }
93
94    /// Unwrap into an Int tensor, or error.
95    pub fn into_int(self) -> Result<Tensor<D, Int>> {
96        match self {
97            Self::Int(t) => Ok(t),
98            other => Err(crate::Error::Msg(format!("expected Int DynTensor, got {:?}", other.dtype()))),
99        }
100    }
101
102    /// Unwrap into a Bool tensor, or error.
103    pub fn into_bool(self) -> Result<Tensor<D, Bool>> {
104        match self {
105            Self::Bool(t) => Ok(t),
106            other => Err(crate::Error::Msg(format!("expected Bool DynTensor, got {:?}", other.dtype()))),
107        }
108    }
109}
110
111// ---- I/O ----
112
113impl<D: Device> DynTensor<D> {
114    /// Create a `DynTensor` from raw little-endian bytes and a runtime [`DType`].
115    ///
116    /// This is the primary entry point for deserialization (safetensors, NPY, …).
117    /// The byte slice length must equal `shape.element_count() * dtype.size_in_bytes()`.
118    pub fn from_bytes(bytes: &[u8], dtype: DType, shape: impl Into<Shape>, device: &D) -> Result<Self> {
119        let shape = shape.into();
120        match dtype.kind() {
121            KindTag::Float => Ok(DynTensor::Float(Tensor::<D, Float>::from_bytes(bytes, shape, (device, dtype.as_float()))?)),
122            KindTag::Int => Ok(DynTensor::Int(Tensor::<D, Int>::from_bytes(bytes, shape, (device, dtype.as_int()))?)),
123            KindTag::Bool => Ok(DynTensor::Bool(Tensor::<D, Bool>::from_bytes(bytes, shape, (device, BoolDType::Bool))?)),
124        }
125    }
126
127    /// Serialize to raw little-endian bytes in logical order.
128    pub fn to_bytes(&self) -> Result<Vec<u8>> {
129        match self {
130            Self::Float(t) => t.to_bytes(),
131            Self::Int(t) => t.to_bytes(),
132            Self::Bool(t) => t.to_bytes(),
133        }
134    }
135}
136
137// ---- From impls ----
138
139impl<D: Device> From<Tensor<D, Float>> for DynTensor<D> {
140    fn from(t: Tensor<D, Float>) -> Self {
141        Self::Float(t)
142    }
143}
144
145impl<D: Device> From<Tensor<D, Int>> for DynTensor<D> {
146    fn from(t: Tensor<D, Int>) -> Self {
147        Self::Int(t)
148    }
149}
150
151impl<D: Device> From<Tensor<D, Bool>> for DynTensor<D> {
152    fn from(t: Tensor<D, Bool>) -> Self {
153        Self::Bool(t)
154    }
155}
156
157// ---- Display ----
158
159impl<D: Device> std::fmt::Display for DynTensor<D> {
160    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
161        match self {
162            Self::Float(t) => write!(f, "{t}"),
163            Self::Int(t) => write!(f, "{t}"),
164            Self::Bool(t) => write!(f, "{t}"),
165        }
166    }
167}
168
169// ---- Clone (Tensor is Clone via Arc) ----
170
171impl<D: Device> Clone for DynTensor<D> {
172    fn clone(&self) -> Self {
173        match self {
174            Self::Float(t) => Self::Float(t.clone()),
175            Self::Int(t) => Self::Int(t.clone()),
176            Self::Bool(t) => Self::Bool(t.clone()),
177        }
178    }
179}
luma_tensor/dynamic/mod.rs

luma_tensor/dynamic/
mod.rs
