Skip to main content

tensor4all_tensorbackend/
logical_tensor.rs

1//! Backend-free logical tensor snapshots for transfer between execution domains.
2
3use num_complex::{Complex32, Complex64};
4use tenferro::{DType, Tensor, TensorRead};
5
6use crate::CpuExecutionContext;
7
8/// Owned column-major scalar payload for a [`LogicalTensor`].
9///
10/// The enum contains values only; it cannot carry an executor, backend,
11/// runtime, cache, pointer, or address identity.
12#[derive(Clone, Debug, PartialEq)]
13pub enum LogicalTensorData {
14    /// 32-bit real values.
15    F32(Vec<f32>),
16    /// 64-bit real values.
17    F64(Vec<f64>),
18    /// 32-bit signed integer values.
19    I32(Vec<i32>),
20    /// 64-bit signed integer values.
21    I64(Vec<i64>),
22    /// Boolean values.
23    Bool(Vec<bool>),
24    /// 32-bit complex values.
25    C32(Vec<Complex32>),
26    /// 64-bit complex values.
27    C64(Vec<Complex64>),
28}
29
30impl LogicalTensorData {
31    /// Return the scalar dtype represented by this payload.
32    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    /// Return the number of logical scalar values.
45    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    /// Return whether the payload contains no values.
58    pub fn is_empty(&self) -> bool {
59        self.len() == 0
60    }
61}
62
63/// Error returned while validating, snapshotting, or reconstructing a logical tensor.
64#[derive(Debug, thiserror::Error)]
65pub enum LogicalTensorError {
66    /// The shape product overflowed `usize`.
67    #[error("logical tensor shape product overflows usize: {shape:?}")]
68    ShapeOverflow {
69        /// Rejected tensor shape.
70        shape: Vec<usize>,
71    },
72    /// The payload length does not match the shape.
73    #[error(
74        "logical tensor element count mismatch: shape {shape:?} requires {expected}, got {actual}"
75    )]
76    ElementCountMismatch {
77        /// Tensor shape.
78        shape: Vec<usize>,
79        /// Required element count.
80        expected: usize,
81        /// Supplied element count.
82        actual: usize,
83    },
84    /// Tenferro could not read or construct the host tensor.
85    #[error("logical tensor {operation} failed: {source}")]
86    Tensor {
87        /// Operation that failed.
88        operation: &'static str,
89        /// Original tenferro diagnostic.
90        #[source]
91        source: tenferro_tensor::Error,
92    },
93}
94
95/// Backend-free, dtype-preserving logical host tensor in column-major order.
96///
97/// Use [`CpuExecutionContext::reconstruct`] in the receiving execution domain.
98/// The adapter, not tensor4all, owns serialization and transport.
99///
100/// # Examples
101///
102/// ```
103/// use tensor4all_tensorbackend::{LogicalTensor, LogicalTensorData};
104///
105/// let tensor = LogicalTensor::new(
106///     vec![2, 2],
107///     LogicalTensorData::F64(vec![1.0, 2.0, 3.0, 4.0]),
108/// )?;
109/// assert_eq!(tensor.shape(), &[2, 2]);
110/// assert_eq!(tensor.data().len(), 4);
111/// # Ok::<(), tensor4all_tensorbackend::LogicalTensorError>(())
112/// ```
113#[derive(Clone, Debug, PartialEq)]
114pub struct LogicalTensor {
115    shape: Vec<usize>,
116    data: LogicalTensorData,
117}
118
119impl LogicalTensor {
120    /// Create a validated column-major logical tensor snapshot.
121    ///
122    /// # Errors
123    ///
124    /// Returns [`LogicalTensorError::ShapeOverflow`] when the shape product
125    /// overflows, or [`LogicalTensorError::ElementCountMismatch`] when the data
126    /// length differs from that product.
127    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    /// Snapshot a native host tensor without retaining execution identity.
145    ///
146    /// # Errors
147    ///
148    /// Returns [`LogicalTensorError::Tensor`] if the native tensor is not
149    /// host-readable with its declared dtype.
150    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    /// Return the logical shape.
220    pub fn shape(&self) -> &[usize] {
221        &self.shape
222    }
223
224    /// Return the scalar dtype.
225    pub fn dtype(&self) -> DType {
226        self.data.dtype()
227    }
228
229    /// Return the owned scalar payload.
230    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    /// Reconstruct a logical host tensor in this receiving execution domain.
267    ///
268    /// The returned value carries only host tensor data. Its first plain, graph,
269    /// or eager operation must use this same explicit context.
270    ///
271    /// # Errors
272    ///
273    /// Returns [`LogicalTensorError`] if tenferro rejects the validated shape or
274    /// dtype payload.
275    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}