tensor4all_tensorbackend/
logical_tensor.rs1use num_complex::{Complex32, Complex64};
4use tenferro::{DType, Tensor, TensorRead};
5
6use crate::CpuExecutionContext;
7
8#[derive(Clone, Debug, PartialEq)]
13pub enum LogicalTensorData {
14 F32(Vec<f32>),
16 F64(Vec<f64>),
18 I32(Vec<i32>),
20 I64(Vec<i64>),
22 Bool(Vec<bool>),
24 C32(Vec<Complex32>),
26 C64(Vec<Complex64>),
28}
29
30impl LogicalTensorData {
31 pub fn dtype(&self) -> DType {
33 match self {
34 Self::F32(_) => DType::F32,
35 Self::F64(_) => DType::F64,
36 Self::I32(_) => DType::I32,
37 Self::I64(_) => DType::I64,
38 Self::Bool(_) => DType::Bool,
39 Self::C32(_) => DType::C32,
40 Self::C64(_) => DType::C64,
41 }
42 }
43
44 pub fn len(&self) -> usize {
46 match self {
47 Self::F32(values) => values.len(),
48 Self::F64(values) => values.len(),
49 Self::I32(values) => values.len(),
50 Self::I64(values) => values.len(),
51 Self::Bool(values) => values.len(),
52 Self::C32(values) => values.len(),
53 Self::C64(values) => values.len(),
54 }
55 }
56
57 pub fn is_empty(&self) -> bool {
59 self.len() == 0
60 }
61}
62
63#[derive(Debug, thiserror::Error)]
65pub enum LogicalTensorError {
66 #[error("logical tensor shape product overflows usize: {shape:?}")]
68 ShapeOverflow {
69 shape: Vec<usize>,
71 },
72 #[error(
74 "logical tensor element count mismatch: shape {shape:?} requires {expected}, got {actual}"
75 )]
76 ElementCountMismatch {
77 shape: Vec<usize>,
79 expected: usize,
81 actual: usize,
83 },
84 #[error("logical tensor {operation} failed: {source}")]
86 Tensor {
87 operation: &'static str,
89 #[source]
91 source: tenferro_tensor::Error,
92 },
93}
94
95#[derive(Clone, Debug, PartialEq)]
114pub struct LogicalTensor {
115 shape: Vec<usize>,
116 data: LogicalTensorData,
117}
118
119impl LogicalTensor {
120 pub fn new(shape: Vec<usize>, data: LogicalTensorData) -> Result<Self, LogicalTensorError> {
128 let expected = shape
129 .iter()
130 .try_fold(1_usize, |count, &dim| count.checked_mul(dim).ok_or(()));
131 let expected = expected.map_err(|()| LogicalTensorError::ShapeOverflow {
132 shape: shape.clone(),
133 })?;
134 if data.len() != expected {
135 return Err(LogicalTensorError::ElementCountMismatch {
136 shape,
137 expected,
138 actual: data.len(),
139 });
140 }
141 Ok(Self { shape, data })
142 }
143
144 pub fn from_native(tensor: &Tensor) -> Result<Self, LogicalTensorError> {
151 let data = match tensor.dtype() {
152 DType::F32 => LogicalTensorData::F32(
153 tensor
154 .as_slice::<f32>()
155 .map_err(|source| LogicalTensorError::Tensor {
156 operation: "snapshot",
157 source,
158 })?
159 .to_vec(),
160 ),
161 DType::F64 => LogicalTensorData::F64(
162 tensor
163 .as_slice::<f64>()
164 .map_err(|source| LogicalTensorError::Tensor {
165 operation: "snapshot",
166 source,
167 })?
168 .to_vec(),
169 ),
170 DType::I32 => LogicalTensorData::I32(
171 tensor
172 .as_slice::<i32>()
173 .map_err(|source| LogicalTensorError::Tensor {
174 operation: "snapshot",
175 source,
176 })?
177 .to_vec(),
178 ),
179 DType::I64 => LogicalTensorData::I64(
180 tensor
181 .as_slice::<i64>()
182 .map_err(|source| LogicalTensorError::Tensor {
183 operation: "snapshot",
184 source,
185 })?
186 .to_vec(),
187 ),
188 DType::Bool => LogicalTensorData::Bool(
189 tensor
190 .as_slice::<bool>()
191 .map_err(|source| LogicalTensorError::Tensor {
192 operation: "snapshot",
193 source,
194 })?
195 .to_vec(),
196 ),
197 DType::C32 => LogicalTensorData::C32(
198 tensor
199 .as_slice::<Complex32>()
200 .map_err(|source| LogicalTensorError::Tensor {
201 operation: "snapshot",
202 source,
203 })?
204 .to_vec(),
205 ),
206 DType::C64 => LogicalTensorData::C64(
207 tensor
208 .as_slice::<Complex64>()
209 .map_err(|source| LogicalTensorError::Tensor {
210 operation: "snapshot",
211 source,
212 })?
213 .to_vec(),
214 ),
215 };
216 Self::new(tensor.shape().to_vec(), data)
217 }
218
219 pub fn shape(&self) -> &[usize] {
221 &self.shape
222 }
223
224 pub fn dtype(&self) -> DType {
226 self.data.dtype()
227 }
228
229 pub fn data(&self) -> &LogicalTensorData {
231 &self.data
232 }
233
234 fn to_native(&self) -> Result<Tensor, LogicalTensorError> {
235 let result = match &self.data {
236 LogicalTensorData::F32(values) => {
237 Tensor::from_vec_col_major(self.shape.clone(), values.clone())
238 }
239 LogicalTensorData::F64(values) => {
240 Tensor::from_vec_col_major(self.shape.clone(), values.clone())
241 }
242 LogicalTensorData::I32(values) => {
243 Tensor::from_vec_col_major(self.shape.clone(), values.clone())
244 }
245 LogicalTensorData::I64(values) => {
246 Tensor::from_vec_col_major(self.shape.clone(), values.clone())
247 }
248 LogicalTensorData::Bool(values) => {
249 Tensor::from_vec_col_major(self.shape.clone(), values.clone())
250 }
251 LogicalTensorData::C32(values) => {
252 Tensor::from_vec_col_major(self.shape.clone(), values.clone())
253 }
254 LogicalTensorData::C64(values) => {
255 Tensor::from_vec_col_major(self.shape.clone(), values.clone())
256 }
257 };
258 result.map_err(|source| LogicalTensorError::Tensor {
259 operation: "reconstruction",
260 source,
261 })
262 }
263}
264
265impl CpuExecutionContext {
266 pub fn reconstruct(&self, tensor: &LogicalTensor) -> Result<Tensor, LogicalTensorError> {
276 let host = tensor.to_native()?;
277 self.with_session(|session| session.upload_host_tensor(TensorRead::from_tensor(&host)))
278 .map_err(|source| LogicalTensorError::Tensor {
279 operation: "target-context upload",
280 source,
281 })
282 }
283}
284
285#[cfg(test)]
286mod tests {
287 use super::*;
288 use tenferro_cpu::CpuBackend;
289
290 #[test]
291 fn round_trip_preserves_all_host_dtypes_in_target_context() {
292 macro_rules! assert_round_trip {
293 ($ty:ty, $variant:ident, $dtype:expr, $values:expr) => {{
294 let values: Vec<$ty> = $values;
295 let native =
296 Tensor::from_vec_col_major(vec![values.len()], values.clone()).unwrap();
297 let logical = LogicalTensor::from_native(&native).unwrap();
298 let target =
299 CpuExecutionContext::from_backend(CpuBackend::with_threads(1).unwrap());
300 let rebuilt = target.reconstruct(&logical).unwrap();
301
302 assert_eq!(logical.shape(), &[values.len()]);
303 assert_eq!(logical.dtype(), $dtype);
304 assert_eq!(logical.data().len(), values.len());
305 assert!(matches!(logical.data(), LogicalTensorData::$variant(_)));
306 assert_eq!(rebuilt.as_slice::<$ty>().unwrap(), values);
307 }};
308 }
309
310 assert_round_trip!(f32, F32, DType::F32, vec![1.0, 2.0]);
311 assert_round_trip!(f64, F64, DType::F64, vec![1.0, 2.0]);
312 assert_round_trip!(i32, I32, DType::I32, vec![1, 2]);
313 assert_round_trip!(i64, I64, DType::I64, vec![1, 2]);
314 assert_round_trip!(bool, Bool, DType::Bool, vec![true, false]);
315 assert_round_trip!(
316 Complex32,
317 C32,
318 DType::C32,
319 vec![Complex32::new(1.0, 2.0), Complex32::new(3.0, 4.0)]
320 );
321 assert_round_trip!(
322 Complex64,
323 C64,
324 DType::C64,
325 vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, 4.0)]
326 );
327
328 let empty = LogicalTensor::new(vec![0], LogicalTensorData::F64(Vec::new())).unwrap();
329 assert!(empty.data().is_empty());
330 }
331
332 #[test]
333 fn validation_rejects_overflow_and_length_mismatch() {
334 assert!(matches!(
335 LogicalTensor::new(vec![usize::MAX, 2], LogicalTensorData::F64(Vec::new())),
336 Err(LogicalTensorError::ShapeOverflow { .. })
337 ));
338 assert!(matches!(
339 LogicalTensor::new(vec![2], LogicalTensorData::F64(vec![1.0])),
340 Err(LogicalTensorError::ElementCountMismatch { .. })
341 ));
342 }
343}