tenshift_core/sample/
tensor.rs

1use super::ops::{cast_numeric_slice, expected_byte_len, numeric_to_le_bytes};
2use crate::error::{Error, Result};
3use std::sync::{Arc, OnceLock};
4
5#[cfg(feature = "uring")]
6use wireshift::{Buffer, Completed};
7
8/// A typed, shaped buffer of data.
9///
10/// Tensors hold raw bytes with a dtype and shape. Data is reference-counted
11/// so cloning a tensor is cheap (no copy). The actual bytes are only copied
12/// when mutated (copy-on-write).
13#[derive(Debug, Clone)]
14pub struct Tensor {
15    /// The raw data.
16    pub(crate) data: Arc<TensorData>,
17    /// Element type.
18    pub(crate) dtype: DType,
19    /// Shape dimensions. Arc-shared to avoid per-clone Vec allocation.
20    pub(crate) shape: Arc<[usize]>,
21}
22
23#[derive(Debug, Default)]
24pub(crate) struct TensorData {
25    pub(crate) bytes: TensorBytes,
26    pub(crate) f32_cache: OnceLock<Vec<f32>>,
27    pub(crate) f64_cache: OnceLock<Vec<f64>>,
28    pub(crate) i32_cache: OnceLock<Vec<i32>>,
29    pub(crate) i64_cache: OnceLock<Vec<i64>>,
30}
31
32#[derive(Debug)]
33pub(crate) enum TensorBytes {
34    Shared(Arc<[u8]>),
35    #[cfg(feature = "uring")]
36    Pooled(Buffer<Completed>),
37}
38
39impl Default for TensorBytes {
40    fn default() -> Self {
41        Self::Shared(Arc::from([]))
42    }
43}
44
45impl TensorBytes {
46    fn as_slice(&self) -> &[u8] {
47        match self {
48            Self::Shared(bytes) => bytes,
49            #[cfg(feature = "uring")]
50            Self::Pooled(buffer) => buffer.filled(),
51        }
52    }
53
54    fn len(&self) -> usize {
55        self.as_slice().len()
56    }
57}
58
59/// Supported element types.
60#[non_exhaustive]
61#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
62pub enum DType {
63    /// 32-bit float.
64    F32,
65    /// 64-bit float.
66    F64,
67    /// 8-bit unsigned integer.
68    U8,
69    /// 32-bit signed integer.
70    I32,
71    /// 64-bit signed integer.
72    I64,
73    /// Raw bytes (untyped).
74    Bytes,
75}
76
77impl DType {
78    /// Size of one element in bytes.
79    pub fn element_size(self) -> usize {
80        match self {
81            Self::F32 | Self::I32 => 4,
82            Self::F64 | Self::I64 => 8,
83            Self::U8 | Self::Bytes => 1,
84        }
85    }
86
87    /// Human-readable name for this dtype.
88    pub fn name(self) -> &'static str {
89        match self {
90            Self::F32 => "f32",
91            Self::F64 => "f64",
92            Self::U8 => "u8",
93            Self::I32 => "i32",
94            Self::I64 => "i64",
95            Self::Bytes => "bytes",
96        }
97    }
98}
99
100impl std::fmt::Display for DType {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        f.write_str(self.name())
103    }
104}
105
106impl Tensor {
107    /// Create a float32 tensor from a slice of `f32` values.
108    pub fn f32(data: &[f32], shape: Vec<usize>) -> Self {
109        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::F32, shape)
110    }
111
112    /// Create a float64 tensor from a slice of `f64` values.
113    pub fn f64(data: &[f64], shape: Vec<usize>) -> Self {
114        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::F64, shape)
115    }
116
117    /// Create a uint8 tensor (images, raw bytes).
118    pub fn u8(data: Vec<u8>, shape: Vec<usize>) -> Self {
119        Self::from_bytes_unchecked(data, DType::U8, shape)
120    }
121
122    /// Create an int32 tensor.
123    pub fn i32(data: &[i32], shape: Vec<usize>) -> Self {
124        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::I32, shape)
125    }
126
127    /// Create an int64 tensor (labels, indices).
128    pub fn i64(data: &[i64], shape: Vec<usize>) -> Self {
129        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::I64, shape)
130    }
131
132    /// Create a raw bytes tensor (untyped).
133    pub fn bytes(data: Vec<u8>) -> Self {
134        let len = data.len();
135        Self::from_bytes_unchecked(data, DType::Bytes, vec![len])
136    }
137
138    /// Create a raw bytes tensor from a completed `wireshift` buffer.
139    #[cfg(feature = "uring")]
140    pub fn bytes_from_completed(buffer: Buffer<Completed>) -> Self {
141        let len = buffer.filled_len();
142        Self::from_storage_unchecked(TensorBytes::Pooled(buffer), DType::Bytes, vec![len])
143    }
144
145    /// Create a tensor from raw bytes, dtype, and shape.
146    ///
147    /// # Errors
148    /// Returns an error if the provided byte length does not match the product of shape elements
149    /// multiplied by the element size for the given `dtype`.
150    pub fn from_bytes(data: Vec<u8>, dtype: DType, shape: Vec<usize>) -> Result<Self> {
151        let expected = expected_byte_len(dtype, &shape)?;
152        if data.len() != expected {
153            return Err(Error::InvalidConfig {
154                reason: format!(
155                    "tensor byte length {} does not match dtype {:?} and shape {:?} (expected {expected})",
156                    data.len(),
157                    dtype,
158                    shape
159                ),
160            });
161        }
162
163        Ok(Self::from_bytes_unchecked(data, dtype, shape))
164    }
165
166    /// The element type of this tensor.
167    pub fn dtype(&self) -> DType {
168        self.dtype
169    }
170
171    /// The shape of this tensor.
172    pub fn shape(&self) -> &[usize] {
173        &self.shape
174    }
175
176    /// Total number of elements.
177    ///
178    /// Returns `None` if the element count overflows `usize`.
179    pub fn num_elements(&self) -> Option<usize> {
180        self.shape
181            .iter()
182            .try_fold(1_usize, |acc, &dim| acc.checked_mul(dim))
183    }
184
185    /// Raw byte data (for zero-copy transfer to numpy/torch).
186    pub fn as_bytes(&self) -> &[u8] {
187        self.data.bytes.as_slice()
188    }
189
190    /// Interpret data as f32 slice.
191    ///
192    /// # Errors
193    /// Returns an error if the tensor's dtype is not `f32`.
194    pub fn try_as_f32(&self) -> Result<&[f32]> {
195        if self.dtype != DType::F32 {
196            return Err(Error::InvalidConfig {
197                reason: format!("expected dtype f32, got {}", self.dtype),
198            });
199        }
200        cast_numeric_slice(
201            self.data.bytes.as_slice(),
202            &self.data.f32_cache,
203            f32::from_le_bytes,
204        )
205    }
206
207    /// Interpret data as f64 slice.
208    ///
209    /// # Errors
210    /// Returns an error if the tensor's dtype is not `f64`.
211    pub fn try_as_f64(&self) -> Result<&[f64]> {
212        if self.dtype != DType::F64 {
213            return Err(Error::InvalidConfig {
214                reason: format!("expected dtype f64, got {}", self.dtype),
215            });
216        }
217        cast_numeric_slice(
218            self.data.bytes.as_slice(),
219            &self.data.f64_cache,
220            f64::from_le_bytes,
221        )
222    }
223
224    /// Interpret data as i32 slice.
225    ///
226    /// # Errors
227    /// Returns an error if the tensor's dtype is not `i32`.
228    pub fn try_as_i32(&self) -> Result<&[i32]> {
229        if self.dtype != DType::I32 {
230            return Err(Error::InvalidConfig {
231                reason: format!("expected dtype i32, got {}", self.dtype),
232            });
233        }
234        cast_numeric_slice(
235            self.data.bytes.as_slice(),
236            &self.data.i32_cache,
237            i32::from_le_bytes,
238        )
239    }
240
241    /// Interpret data as i64 slice.
242    ///
243    /// # Errors
244    /// Returns an error if the tensor's dtype is not `i64`.
245    pub fn try_as_i64(&self) -> Result<&[i64]> {
246        if self.dtype != DType::I64 {
247            return Err(Error::InvalidConfig {
248                reason: format!("expected dtype i64, got {}", self.dtype),
249            });
250        }
251        cast_numeric_slice(
252            self.data.bytes.as_slice(),
253            &self.data.i64_cache,
254            i64::from_le_bytes,
255        )
256    }
257
258    /// Byte length of the data buffer.
259    pub fn byte_len(&self) -> usize {
260        self.data.bytes.len()
261    }
262
263    pub(crate) fn from_bytes_unchecked(data: Vec<u8>, dtype: DType, shape: Vec<usize>) -> Self {
264        Self::from_storage_unchecked(TensorBytes::Shared(Arc::from(data)), dtype, shape)
265    }
266
267    pub(crate) fn from_storage_unchecked(
268        data: TensorBytes,
269        dtype: DType,
270        shape: Vec<usize>,
271    ) -> Self {
272        Self {
273            data: Arc::new(TensorData {
274                bytes: data,
275                f32_cache: OnceLock::new(),
276                f64_cache: OnceLock::new(),
277                i32_cache: OnceLock::new(),
278                i64_cache: OnceLock::new(),
279            }),
280            dtype,
281            shape: Arc::from(shape),
282        }
283    }
284}
tenshift_core/sample/tensor.rs

tenshift_core/sample/
tensor.rs
