spatialrust_tensor/
lib.rs

1//! Small, runtime-independent tensor descriptors and CPU storage views.
2//!
3//! Shape, element strides, byte offset, device, and ownership are explicit.
4//! This crate never uploads, downloads, or otherwise migrates storage. DLPack
5//! FFI and image bridges are additive features built on this data model.
6
7#![deny(unsafe_code)]
8#![warn(missing_docs)]
9
10use std::{any::Any, fmt::Debug, ops::Range, sync::Arc};
11
12#[cfg(feature = "image")]
13mod image;
14#[cfg(feature = "image")]
15pub use image::{
16    interleaved_image_view, pack_interleaved_image, pack_planar_image, planar_image_view,
17    TensorElement,
18};
19#[cfg(feature = "spatial")]
20mod spatial;
21#[cfg(feature = "spatial")]
22pub use spatial::{spatial_f32_field_view, SpatialTensorBridgeError};
23#[cfg(feature = "dlpack")]
24#[allow(unsafe_code)]
25mod dlpack;
26#[cfg(feature = "dlpack")]
27pub use dlpack::{
28    release_dlpack_legacy_raw, release_dlpack_raw, DlpackError, DlpackExport, DlpackImport,
29    DlpackLegacyExport, DLPACK_MAJOR, DLPACK_MINOR,
30};
31
32/// Tensor construction and layout errors.
33#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
34pub enum TensorError {
35    /// The data type has no bits or lanes, or cannot address whole bytes.
36    #[error("invalid data type: bits and lanes must be non-zero and form whole bytes")]
37    InvalidDataType,
38    /// The stride rank differs from the shape rank.
39    #[error("stride rank {strides} differs from shape rank {shape}")]
40    RankMismatch {
41        /// Number of shape dimensions.
42        shape: usize,
43        /// Number of strides.
44        strides: usize,
45    },
46    /// Shape, stride, or byte-range arithmetic overflowed.
47    #[error("tensor layout overflows addressable memory")]
48    LayoutOverflow,
49    /// The byte offset or reachable tensor span lies outside storage.
50    #[error("tensor byte range {start}..{end} lies outside {available} bytes")]
51    StorageOutOfBounds {
52        /// First reachable byte.
53        start: i128,
54        /// Exclusive final reachable byte.
55        end: i128,
56        /// Available storage size.
57        available: usize,
58    },
59    /// A host slice was paired with a device that is not host-accessible.
60    #[error("device {0:?} is not directly host-accessible")]
61    DeviceNotHostAccessible(Device),
62    /// Typed storage does not match its descriptor dtype.
63    #[error("typed storage requires {expected:?}, descriptor declares {actual:?}")]
64    DataTypeMismatch {
65        /// Storage dtype.
66        expected: DataType,
67        /// Descriptor dtype.
68        actual: DataType,
69    },
70}
71
72/// DLPack-compatible scalar category without depending on a runtime header.
73#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
74#[repr(u8)]
75pub enum DataTypeCode {
76    /// Signed integer.
77    Int = 0,
78    /// Unsigned integer.
79    UInt = 1,
80    /// IEEE floating point.
81    Float = 2,
82    /// Brain floating point.
83    BFloat = 4,
84    /// Complex floating point.
85    Complex = 5,
86    /// Boolean value.
87    Bool = 6,
88}
89
90/// Scalar category, bit width, and vector lane count.
91#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
92pub struct DataType {
93    code: DataTypeCode,
94    bits: u8,
95    lanes: u16,
96}
97
98impl DataType {
99    /// Unsigned 8-bit scalar.
100    pub const U8: Self = Self::new_unchecked(DataTypeCode::UInt, 8, 1);
101    /// Unsigned 16-bit scalar.
102    pub const U16: Self = Self::new_unchecked(DataTypeCode::UInt, 16, 1);
103    /// Unsigned 32-bit scalar.
104    pub const U32: Self = Self::new_unchecked(DataTypeCode::UInt, 32, 1);
105    /// Signed 8-bit scalar.
106    pub const I8: Self = Self::new_unchecked(DataTypeCode::Int, 8, 1);
107    /// Signed 16-bit scalar.
108    pub const I16: Self = Self::new_unchecked(DataTypeCode::Int, 16, 1);
109    /// Signed 32-bit scalar.
110    pub const I32: Self = Self::new_unchecked(DataTypeCode::Int, 32, 1);
111    /// Signed 64-bit scalar.
112    pub const I64: Self = Self::new_unchecked(DataTypeCode::Int, 64, 1);
113    /// IEEE binary16 scalar.
114    pub const F16: Self = Self::new_unchecked(DataTypeCode::Float, 16, 1);
115    /// Brain floating-point 16-bit scalar.
116    pub const BF16: Self = Self::new_unchecked(DataTypeCode::BFloat, 16, 1);
117    /// IEEE binary32 scalar.
118    pub const F32: Self = Self::new_unchecked(DataTypeCode::Float, 32, 1);
119    /// IEEE binary64 scalar.
120    pub const F64: Self = Self::new_unchecked(DataTypeCode::Float, 64, 1);
121    /// Eight-bit boolean scalar, matching common DLPack producers.
122    pub const BOOL: Self = Self::new_unchecked(DataTypeCode::Bool, 8, 1);
123
124    const fn new_unchecked(code: DataTypeCode, bits: u8, lanes: u16) -> Self {
125        Self { code, bits, lanes }
126    }
127
128    /// Creates a byte-addressable scalar or vector data type.
129    pub fn try_new(code: DataTypeCode, bits: u8, lanes: u16) -> Result<Self, TensorError> {
130        let total_bits = usize::from(bits)
131            .checked_mul(usize::from(lanes))
132            .ok_or(TensorError::InvalidDataType)?;
133        if bits == 0 || lanes == 0 || total_bits % 8 != 0 {
134            return Err(TensorError::InvalidDataType);
135        }
136        Ok(Self { code, bits, lanes })
137    }
138
139    /// Returns the scalar category.
140    pub const fn code(self) -> DataTypeCode {
141        self.code
142    }
143
144    /// Returns bits per lane.
145    pub const fn bits(self) -> u8 {
146        self.bits
147    }
148
149    /// Returns the vector lane count.
150    pub const fn lanes(self) -> u16 {
151        self.lanes
152    }
153
154    /// Returns bytes occupied by one tensor element.
155    pub fn element_size(self) -> usize {
156        usize::from(self.bits) * usize::from(self.lanes) / 8
157    }
158}
159
160/// Device category represented independently of any execution backend.
161#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
162pub enum DeviceKind {
163    /// Ordinary CPU memory.
164    Cpu,
165    /// CUDA device memory.
166    Cuda,
167    /// CUDA-pinned host memory.
168    CudaHost,
169    /// OpenCL device memory.
170    OpenCl,
171    /// Vulkan device memory.
172    Vulkan,
173    /// Metal device memory.
174    Metal,
175    /// Verilog simulator buffer.
176    Vpi,
177    /// ROCm device memory.
178    Rocm,
179    /// ROCm-pinned host memory.
180    RocmHost,
181    /// Backend-specific external memory.
182    External,
183    /// CUDA managed memory.
184    CudaManaged,
185    /// oneAPI device memory.
186    OneApi,
187    /// WebGPU device memory.
188    WebGpu,
189    /// Hexagon device memory.
190    Hexagon,
191    /// Microsoft MAIA device memory.
192    Maia,
193    /// AWS Trainium device memory.
194    Trainium,
195    /// Google TPU device memory.
196    Tpu,
197    /// Google TPU pinned host memory.
198    TpuHost,
199}
200
201/// Device category and backend-local ordinal.
202#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
203pub struct Device {
204    /// Device category.
205    pub kind: DeviceKind,
206    /// Backend-local device ordinal.
207    pub id: i32,
208}
209
210impl Device {
211    /// Main CPU device.
212    pub const CPU: Self = Self { kind: DeviceKind::Cpu, id: 0 };
213
214    /// Returns whether Rust may safely expose this memory as a host byte slice.
215    pub const fn is_host_accessible(self) -> bool {
216        matches!(
217            self.kind,
218            DeviceKind::Cpu | DeviceKind::CudaHost | DeviceKind::RocmHost | DeviceKind::TpuHost
219        )
220    }
221}
222
223/// Whether a tensor object owns or borrows its backing allocation.
224#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
225pub enum Ownership {
226    /// The object owns and drops its storage.
227    Owned,
228    /// The object cannot outlive borrowed storage.
229    Borrowed,
230}
231
232/// Shape, element strides, type, offset, and device for one tensor.
233#[derive(Clone, Debug, PartialEq, Eq)]
234pub struct TensorDescriptor {
235    dtype: DataType,
236    shape: Vec<usize>,
237    strides: Option<Vec<isize>>,
238    byte_offset: usize,
239    device: Device,
240}
241
242impl TensorDescriptor {
243    /// Creates a compact C-order tensor on a device.
244    pub fn contiguous(dtype: DataType, shape: Vec<usize>, device: Device) -> Self {
245        Self { dtype, shape, strides: None, byte_offset: 0, device }
246    }
247
248    /// Creates an explicitly strided tensor. Strides are measured in elements.
249    pub fn try_strided(
250        dtype: DataType,
251        shape: Vec<usize>,
252        strides: Vec<isize>,
253        byte_offset: usize,
254        device: Device,
255    ) -> Result<Self, TensorError> {
256        if shape.len() != strides.len() {
257            return Err(TensorError::RankMismatch { shape: shape.len(), strides: strides.len() });
258        }
259        let descriptor = Self { dtype, shape, strides: Some(strides), byte_offset, device };
260        descriptor.required_byte_range()?;
261        Ok(descriptor)
262    }
263
264    /// Returns the data type.
265    pub const fn dtype(&self) -> DataType {
266        self.dtype
267    }
268
269    /// Returns dimensions in logical axis order.
270    pub fn shape(&self) -> &[usize] {
271        &self.shape
272    }
273
274    /// Returns element strides, or `None` for compact C order.
275    pub fn strides(&self) -> Option<&[isize]> {
276        self.strides.as_deref()
277    }
278
279    /// Returns the byte offset from the allocation start to logical element zero.
280    pub const fn byte_offset(&self) -> usize {
281        self.byte_offset
282    }
283
284    /// Returns the storage device.
285    pub const fn device(&self) -> Device {
286        self.device
287    }
288
289    /// Returns the number of logical elements, including one for a scalar.
290    pub fn element_count(&self) -> Result<usize, TensorError> {
291        self.shape
292            .iter()
293            .try_fold(1usize, |count, &dimension| count.checked_mul(dimension))
294            .ok_or(TensorError::LayoutOverflow)
295    }
296
297    /// Returns whether the tensor is compact in C order.
298    pub fn is_c_contiguous(&self) -> bool {
299        let Some(strides) = &self.strides else { return true };
300        let mut expected = 1usize;
301        for (&dimension, &stride) in self.shape.iter().zip(strides).rev() {
302            if dimension > 1 && usize::try_from(stride).ok() != Some(expected) {
303                return false;
304            }
305            let Some(next) = expected.checked_mul(dimension.max(1)) else { return false };
306            expected = next;
307        }
308        true
309    }
310
311    /// Computes the smallest allocation byte range reachable by the layout.
312    pub fn required_byte_range(&self) -> Result<Range<usize>, TensorError> {
313        if self.element_count()? == 0 {
314            return Ok(self.byte_offset..self.byte_offset);
315        }
316        let item_size =
317            i128::try_from(self.dtype.element_size()).map_err(|_| TensorError::LayoutOverflow)?;
318        let origin = i128::try_from(self.byte_offset).map_err(|_| TensorError::LayoutOverflow)?;
319        let mut minimum = origin;
320        let mut maximum = origin;
321        if let Some(strides) = &self.strides {
322            for (&dimension, &stride) in self.shape.iter().zip(strides) {
323                let steps = i128::try_from(dimension.saturating_sub(1))
324                    .map_err(|_| TensorError::LayoutOverflow)?;
325                let stride_bytes =
326                    (stride as i128).checked_mul(item_size).ok_or(TensorError::LayoutOverflow)?;
327                let delta = steps.checked_mul(stride_bytes).ok_or(TensorError::LayoutOverflow)?;
328                if delta < 0 {
329                    minimum = minimum.checked_add(delta).ok_or(TensorError::LayoutOverflow)?;
330                } else {
331                    maximum = maximum.checked_add(delta).ok_or(TensorError::LayoutOverflow)?;
332                }
333            }
334        } else {
335            let count =
336                i128::try_from(self.element_count()?).map_err(|_| TensorError::LayoutOverflow)?;
337            maximum = maximum
338                .checked_add((count - 1).checked_mul(item_size).ok_or(TensorError::LayoutOverflow)?)
339                .ok_or(TensorError::LayoutOverflow)?;
340        }
341        let end = maximum.checked_add(item_size).ok_or(TensorError::LayoutOverflow)?;
342        let start = usize::try_from(minimum).map_err(|_| TensorError::StorageOutOfBounds {
343            start: minimum,
344            end,
345            available: 0,
346        })?;
347        let end = usize::try_from(end).map_err(|_| TensorError::LayoutOverflow)?;
348        Ok(start..end)
349    }
350
351    fn validate_storage(&self, available: usize) -> Result<(), TensorError> {
352        if !self.device.is_host_accessible() {
353            return Err(TensorError::DeviceNotHostAccessible(self.device));
354        }
355        let range = self.required_byte_range()?;
356        if range.start > available || range.end > available {
357            return Err(TensorError::StorageOutOfBounds {
358                start: range.start as i128,
359                end: range.end as i128,
360                available,
361            });
362        }
363        Ok(())
364    }
365}
366
367/// Lifetime-bound, zero-copy view of host-accessible tensor storage.
368#[derive(Clone, Debug)]
369pub struct TensorView<'a> {
370    bytes: &'a [u8],
371    descriptor: TensorDescriptor,
372}
373
374impl<'a> TensorView<'a> {
375    /// Validates and borrows a host allocation without copying it.
376    pub fn try_new(bytes: &'a [u8], descriptor: TensorDescriptor) -> Result<Self, TensorError> {
377        descriptor.validate_storage(bytes.len())?;
378        Ok(Self { bytes, descriptor })
379    }
380
381    /// Returns the complete borrowed allocation slice.
382    pub const fn allocation_bytes(&self) -> &'a [u8] {
383        self.bytes
384    }
385
386    /// Returns tensor metadata.
387    pub const fn descriptor(&self) -> &TensorDescriptor {
388        &self.descriptor
389    }
390
391    /// Reports borrowed ownership.
392    pub const fn ownership(&self) -> Ownership {
393        Ownership::Borrowed
394    }
395
396    /// Performs an explicit host-to-host copy retaining the same layout.
397    pub fn to_owned_copy(&self) -> TensorBuffer {
398        TensorBuffer {
399            storage: copy_storage(self.bytes, self.descriptor.dtype()),
400            descriptor: self.descriptor.clone(),
401        }
402    }
403}
404
405/// Owned host allocation paired with a validated tensor descriptor.
406#[derive(Clone, Debug)]
407pub struct TensorBuffer {
408    storage: TensorStorage,
409    descriptor: TensorDescriptor,
410}
411
412#[derive(Clone, Debug)]
413pub(crate) enum TensorStorage {
414    Bytes(Arc<[u8]>),
415    U16(Arc<[u16]>),
416    I16(Arc<[i16]>),
417    U32(Arc<[u32]>),
418    I32(Arc<[i32]>),
419    I64(Arc<[i64]>),
420    F32(Arc<[f32]>),
421    F64(Arc<[f64]>),
422    External(Arc<dyn HostTensorStorage>),
423}
424
425/// Runtime-owned, host-accessible storage retained without a runtime dependency.
426///
427/// Backend crates use this boundary to keep an allocator or runtime value alive
428/// while exposing its stable CPU allocation. Implementations must keep the
429/// returned allocation address and length unchanged for their entire lifetime.
430pub trait HostTensorStorage: Any + Debug + Send + Sync {
431    /// Returns the exact tensor element type stored by this allocation.
432    fn dtype(&self) -> DataType;
433
434    /// Returns the complete host-accessible allocation.
435    fn allocation_bytes(&self) -> &[u8];
436
437    /// Supports backend-specific zero-copy reuse through checked downcasting.
438    fn as_any(&self) -> &dyn Any;
439}
440
441impl TensorStorage {
442    pub(crate) fn as_bytes(&self) -> &[u8] {
443        match self {
444            Self::Bytes(values) => values,
445            Self::U16(values) => bytemuck::cast_slice(values),
446            Self::I16(values) => bytemuck::cast_slice(values),
447            Self::U32(values) => bytemuck::cast_slice(values),
448            Self::I32(values) => bytemuck::cast_slice(values),
449            Self::I64(values) => bytemuck::cast_slice(values),
450            Self::F32(values) => bytemuck::cast_slice(values),
451            Self::F64(values) => bytemuck::cast_slice(values),
452            Self::External(storage) => storage.allocation_bytes(),
453        }
454    }
455
456    #[cfg(feature = "dlpack")]
457    pub(crate) fn is_empty(&self) -> bool {
458        self.as_bytes().is_empty()
459    }
460
461    #[cfg(feature = "dlpack")]
462    pub(crate) fn as_ptr(&self) -> *const u8 {
463        self.as_bytes().as_ptr()
464    }
465
466    #[cfg(all(feature = "dlpack", test))]
467    pub(crate) fn strong_count(&self) -> usize {
468        match self {
469            Self::Bytes(values) => Arc::strong_count(values),
470            Self::U16(values) => Arc::strong_count(values),
471            Self::I16(values) => Arc::strong_count(values),
472            Self::U32(values) => Arc::strong_count(values),
473            Self::I32(values) => Arc::strong_count(values),
474            Self::I64(values) => Arc::strong_count(values),
475            Self::F32(values) => Arc::strong_count(values),
476            Self::F64(values) => Arc::strong_count(values),
477            Self::External(storage) => Arc::strong_count(storage),
478        }
479    }
480}
481
482fn copy_storage(bytes: &[u8], dtype: DataType) -> TensorStorage {
483    macro_rules! aligned_copy {
484        ($type:ty, $variant:ident) => {
485            bytemuck::try_cast_slice::<u8, $type>(bytes)
486                .map(|values| TensorStorage::$variant(Arc::from(values)))
487                .unwrap_or_else(|_| TensorStorage::Bytes(Arc::from(bytes)))
488        };
489    }
490    match dtype {
491        DataType::U16 | DataType::F16 | DataType::BF16 => aligned_copy!(u16, U16),
492        DataType::I16 => aligned_copy!(i16, I16),
493        DataType::U32 => aligned_copy!(u32, U32),
494        DataType::I32 => aligned_copy!(i32, I32),
495        DataType::I64 => aligned_copy!(i64, I64),
496        DataType::F32 => aligned_copy!(f32, F32),
497        DataType::F64 => aligned_copy!(f64, F64),
498        _ => TensorStorage::Bytes(Arc::from(bytes)),
499    }
500}
501
502impl PartialEq for TensorBuffer {
503    fn eq(&self, other: &Self) -> bool {
504        self.descriptor == other.descriptor && self.allocation_bytes() == other.allocation_bytes()
505    }
506}
507
508impl Eq for TensorBuffer {}
509
510impl TensorBuffer {
511    /// Validates and takes ownership of a host allocation.
512    pub fn try_new(bytes: Vec<u8>, descriptor: TensorDescriptor) -> Result<Self, TensorError> {
513        descriptor.validate_storage(bytes.len())?;
514        Ok(Self { storage: TensorStorage::Bytes(Arc::from(bytes)), descriptor })
515    }
516
517    /// Validates and owns aligned u16 storage without byte repacking.
518    pub fn try_from_u16(
519        values: Vec<u16>,
520        descriptor: TensorDescriptor,
521    ) -> Result<Self, TensorError> {
522        Self::try_from_storage(TensorStorage::U16(Arc::from(values)), DataType::U16, descriptor)
523    }
524
525    /// Validates and owns aligned i16 storage without byte repacking.
526    pub fn try_from_i16(
527        values: Vec<i16>,
528        descriptor: TensorDescriptor,
529    ) -> Result<Self, TensorError> {
530        Self::try_from_storage(TensorStorage::I16(Arc::from(values)), DataType::I16, descriptor)
531    }
532
533    /// Validates and owns aligned u32 storage without byte repacking.
534    pub fn try_from_u32(
535        values: Vec<u32>,
536        descriptor: TensorDescriptor,
537    ) -> Result<Self, TensorError> {
538        Self::try_from_storage(TensorStorage::U32(Arc::from(values)), DataType::U32, descriptor)
539    }
540
541    /// Validates and owns aligned i32 storage without byte repacking.
542    pub fn try_from_i32(
543        values: Vec<i32>,
544        descriptor: TensorDescriptor,
545    ) -> Result<Self, TensorError> {
546        Self::try_from_storage(TensorStorage::I32(Arc::from(values)), DataType::I32, descriptor)
547    }
548
549    /// Validates and owns aligned i64 storage without byte repacking.
550    pub fn try_from_i64(
551        values: Vec<i64>,
552        descriptor: TensorDescriptor,
553    ) -> Result<Self, TensorError> {
554        Self::try_from_storage(TensorStorage::I64(Arc::from(values)), DataType::I64, descriptor)
555    }
556
557    /// Validates and owns aligned f32 storage without byte repacking.
558    pub fn try_from_f32(
559        values: Vec<f32>,
560        descriptor: TensorDescriptor,
561    ) -> Result<Self, TensorError> {
562        Self::try_from_storage(TensorStorage::F32(Arc::from(values)), DataType::F32, descriptor)
563    }
564
565    /// Validates and owns aligned f64 storage without byte repacking.
566    pub fn try_from_f64(
567        values: Vec<f64>,
568        descriptor: TensorDescriptor,
569    ) -> Result<Self, TensorError> {
570        Self::try_from_storage(TensorStorage::F64(Arc::from(values)), DataType::F64, descriptor)
571    }
572
573    /// Validates and retains a runtime-owned host allocation without copying it.
574    pub fn try_from_host_storage(
575        storage: Arc<dyn HostTensorStorage>,
576        descriptor: TensorDescriptor,
577    ) -> Result<Self, TensorError> {
578        let dtype = storage.dtype();
579        Self::try_from_storage(TensorStorage::External(storage), dtype, descriptor)
580    }
581
582    fn try_from_storage(
583        storage: TensorStorage,
584        expected: DataType,
585        descriptor: TensorDescriptor,
586    ) -> Result<Self, TensorError> {
587        if descriptor.dtype() != expected {
588            return Err(TensorError::DataTypeMismatch { expected, actual: descriptor.dtype() });
589        }
590        descriptor.validate_storage(storage.as_bytes().len())?;
591        Ok(Self { storage, descriptor })
592    }
593
594    /// Returns a zero-copy borrowed view.
595    pub fn view(&self) -> TensorView<'_> {
596        TensorView { bytes: self.storage.as_bytes(), descriptor: self.descriptor.clone() }
597    }
598
599    /// Returns tensor metadata.
600    pub const fn descriptor(&self) -> &TensorDescriptor {
601        &self.descriptor
602    }
603
604    /// Returns the complete owned allocation bytes.
605    pub fn allocation_bytes(&self) -> &[u8] {
606        self.storage.as_bytes()
607    }
608
609    /// Performs an explicit host-to-host copy while preserving typed alignment.
610    pub fn to_owned_copy(&self) -> Self {
611        Self {
612            storage: copy_storage(self.storage.as_bytes(), self.descriptor.dtype()),
613            descriptor: self.descriptor.clone(),
614        }
615    }
616
617    #[cfg(feature = "dlpack")]
618    pub(crate) fn shared_allocation(&self) -> TensorStorage {
619        self.storage.clone()
620    }
621
622    /// Returns shared aligned f32 storage when constructed with [`Self::try_from_f32`].
623    pub fn shared_f32(&self) -> Option<Arc<[f32]>> {
624        match &self.storage {
625            TensorStorage::F32(values) => Some(Arc::clone(values)),
626            _ => None,
627        }
628    }
629
630    /// Returns shared byte storage when constructed with [`Self::try_new`].
631    ///
632    /// This is suitable for zero-copy `u8` and `i8` tensor adapters. Multi-byte
633    /// element types must use their matching typed constructor and accessor.
634    pub fn shared_bytes(&self) -> Option<Arc<[u8]>> {
635        match &self.storage {
636            TensorStorage::Bytes(values) => Some(Arc::clone(values)),
637            _ => None,
638        }
639    }
640
641    /// Returns shared aligned f64 storage when constructed with [`Self::try_from_f64`].
642    pub fn shared_f64(&self) -> Option<Arc<[f64]>> {
643        match &self.storage {
644            TensorStorage::F64(values) => Some(Arc::clone(values)),
645            _ => None,
646        }
647    }
648
649    /// Returns shared aligned i16 storage when constructed with [`Self::try_from_i16`].
650    pub fn shared_i16(&self) -> Option<Arc<[i16]>> {
651        match &self.storage {
652            TensorStorage::I16(values) => Some(Arc::clone(values)),
653            _ => None,
654        }
655    }
656
657    /// Returns shared aligned i32 storage when constructed with [`Self::try_from_i32`].
658    pub fn shared_i32(&self) -> Option<Arc<[i32]>> {
659        match &self.storage {
660            TensorStorage::I32(values) => Some(Arc::clone(values)),
661            _ => None,
662        }
663    }
664
665    /// Returns shared aligned i64 storage when constructed with [`Self::try_from_i64`].
666    pub fn shared_i64(&self) -> Option<Arc<[i64]>> {
667        match &self.storage {
668            TensorStorage::I64(values) => Some(Arc::clone(values)),
669            _ => None,
670        }
671    }
672
673    /// Returns shared aligned u16 storage when constructed with [`Self::try_from_u16`].
674    pub fn shared_u16(&self) -> Option<Arc<[u16]>> {
675        match &self.storage {
676            TensorStorage::U16(values) => Some(Arc::clone(values)),
677            _ => None,
678        }
679    }
680
681    /// Returns shared aligned u32 storage when constructed with [`Self::try_from_u32`].
682    pub fn shared_u32(&self) -> Option<Arc<[u32]>> {
683        match &self.storage {
684            TensorStorage::U32(values) => Some(Arc::clone(values)),
685            _ => None,
686        }
687    }
688
689    /// Returns runtime-owned host storage, when this tensor wraps one.
690    pub fn host_storage(&self) -> Option<&Arc<dyn HostTensorStorage>> {
691        match &self.storage {
692            TensorStorage::External(storage) => Some(storage),
693            _ => None,
694        }
695    }
696
697    /// Reports owned storage.
698    pub const fn ownership(&self) -> Ownership {
699        Ownership::Owned
700    }
701
702    /// Explicitly copies the allocation into bytes and returns its metadata.
703    pub fn into_allocation_bytes_copy(self) -> (Vec<u8>, TensorDescriptor) {
704        (self.storage.as_bytes().to_vec(), self.descriptor)
705    }
706}
707
708#[cfg(test)]
709mod tests {
710    use super::{
711        DataType, Device, DeviceKind, TensorBuffer, TensorDescriptor, TensorError, TensorView,
712    };
713
714    #[test]
715    fn compact_tensor_validates_exact_storage() {
716        let descriptor = TensorDescriptor::contiguous(DataType::F32, vec![2, 3], Device::CPU);
717        let tensor = TensorBuffer::try_new(vec![0; 24], descriptor).unwrap();
718        assert_eq!(tensor.descriptor().element_count().unwrap(), 6);
719        assert!(tensor.descriptor().is_c_contiguous());
720        assert_eq!(tensor.descriptor().required_byte_range().unwrap(), 0..24);
721    }
722
723    #[test]
724    fn strided_roi_and_negative_stride_have_checked_spans() {
725        let roi =
726            TensorDescriptor::try_strided(DataType::U8, vec![2, 3], vec![5, 1], 6, Device::CPU)
727                .unwrap();
728        assert_eq!(roi.required_byte_range().unwrap(), 6..14);
729        TensorView::try_new(&[0; 14], roi).unwrap();
730
731        let reversed =
732            TensorDescriptor::try_strided(DataType::U8, vec![4], vec![-1], 3, Device::CPU).unwrap();
733        assert_eq!(reversed.required_byte_range().unwrap(), 0..4);
734        TensorView::try_new(&[0; 4], reversed).unwrap();
735    }
736
737    #[test]
738    fn rejects_short_storage_and_device_memory_as_host_slice() {
739        let descriptor = TensorDescriptor::contiguous(DataType::U16, vec![4], Device::CPU);
740        assert!(matches!(
741            TensorView::try_new(&[0; 7], descriptor),
742            Err(TensorError::StorageOutOfBounds { .. })
743        ));
744        let cuda = TensorDescriptor::contiguous(
745            DataType::U8,
746            vec![1],
747            Device { kind: DeviceKind::Cuda, id: 0 },
748        );
749        assert!(matches!(
750            TensorView::try_new(&[0], cuda),
751            Err(TensorError::DeviceNotHostAccessible(_))
752        ));
753    }
754
755    #[test]
756    fn zero_sized_and_scalar_shapes_are_distinct() {
757        let empty = TensorDescriptor::contiguous(DataType::U8, vec![2, 0, 3], Device::CPU);
758        assert_eq!(empty.element_count().unwrap(), 0);
759        assert_eq!(empty.required_byte_range().unwrap(), 0..0);
760        let scalar = TensorDescriptor::contiguous(DataType::F64, vec![], Device::CPU);
761        assert_eq!(scalar.element_count().unwrap(), 1);
762        assert_eq!(scalar.required_byte_range().unwrap(), 0..8);
763    }
764}
spatialrust_tensor/lib.rs

spatialrust_tensor/
lib.rs
