lumen_core/base/
tensor.rs

1use std::{borrow::Borrow, hash::Hash, sync::Arc};
2use crate::{AutogradInfo, AutogradMetaT, DType, FloatDType, NumDType, Op, WithDType};
3use super::{Dim, DimCoordinates, DimNCoordinates, Layout, Shape, Storage, StorageArc, StorageIndices};
4
5#[derive(Clone)]
6pub struct Tensor<T: WithDType>(pub(crate) Arc<TensorImpl<T>>);
7
8#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
9pub struct TensorId(usize);
10
11pub(crate) struct TensorImpl<T: WithDType> {
12    pub(crate) id: TensorId,
13    pub(crate) storage: Option<StorageArc<T>>,
14    pub(crate) layout: Layout,
15    pub(crate) meta: T::AutogradMeta,
16}
17
18impl TensorId {
19    pub fn new() -> Self {
20        use std::sync::atomic;
21        static COUNTER: atomic::AtomicUsize = atomic::AtomicUsize::new(1);
22        Self(COUNTER.fetch_add(1, atomic::Ordering::Relaxed))
23    }
24
25    pub fn value(&self) -> usize {
26        self.0
27    }
28}
29
30impl Borrow<usize> for TensorId {
31    fn borrow(&self) -> &usize {
32        &self.0
33    }
34}
35
36impl<T: WithDType> Hash for Tensor<T> {
37    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
38        self.0.id.0.hash(state);
39    }
40} 
41
42impl<T: WithDType> PartialEq for Tensor<T> {
43    fn eq(&self, other: &Self) -> bool {
44        self.0.id.0 == other.0.id.0
45    }
46}
47
48impl<T: WithDType> Eq for Tensor<T> {}
49
50impl<T: WithDType> Tensor<T> {
51    pub fn is_scalar(&self) -> bool {
52        self.shape().is_scalar()
53    }
54
55    pub fn check_scalar(&self) -> crate::Result<()> {
56        if !self.is_scalar() {
57            Err(crate::Error::NotScalar)?
58        } else {
59            Ok(())
60        }
61    }
62
63    pub fn to_scalar(&self) -> crate::Result<T> {
64        self.check_scalar()?;
65        let v = self.storage_read()?.get_unchecked(self.layout().start_offset());
66        Ok(v)
67    }
68
69    pub fn set_scalar(&self, val: T) -> crate::Result<()> {
70        self.check_scalar()?;
71        self.storage_write()?.set_unchecked(self.layout().start_offset(), val);
72        Ok(())
73    }
74    
75    pub fn storage_ptr(&self, start_offset: usize) -> crate::Result<*mut T> {
76        self.0.storage.as_ref()
77            .ok_or(crate::Error::MetaTensor)
78            .map(|s| s.get_ptr(start_offset))
79    }
80
81    pub fn is_meta(&self) -> bool {
82        self.0.storage.is_none()
83    }
84
85    #[allow(unused)]
86    pub(crate) fn storage_clone(&self) -> crate::Result<StorageArc<T>> {
87        self.0.storage.as_ref()
88            .ok_or(crate::Error::MetaTensor)
89            .cloned()
90    }
91}
92
93impl<T: WithDType> Tensor<T> {
94    pub fn id(&self) -> TensorId {
95        self.0.id
96    }
97
98    pub fn shape(&self) -> &Shape {
99        self.0.layout.shape()
100    }
101
102    pub fn dtype(&self) -> DType {
103        T::DTYPE
104    }
105
106    pub fn layout(&self) -> &Layout {
107        &self.0.layout
108    }
109
110    pub fn dims(&self) -> &[usize] {
111        self.shape().dims()
112    }
113
114    pub fn dim<D: Dim>(&self, dim: D) -> crate::Result<usize> {
115        let dim = dim.to_index(self.shape(), "dim")?;
116        Ok(self.dims()[dim])
117    }
118
119    pub fn storage_read(&self) -> crate::Result<std::sync::RwLockReadGuard<'_, Storage<T>>> {
120        self.0.storage.as_ref()
121            .ok_or(crate::Error::MetaTensor)
122            .map(|s| s.read())
123    }
124
125    pub fn storage_write(&self) -> crate::Result<std::sync::RwLockWriteGuard<'_, Storage<T>>> {
126        self.0.storage.as_ref()
127            .ok_or(crate::Error::MetaTensor)
128            .map(|s| s.write())
129    }
130    
131    pub fn element_count(&self) -> usize {
132        self.shape().element_count()
133    }
134
135    pub fn is_contiguous(&self) -> bool {
136        self.layout().is_contiguous()
137    }
138
139    pub fn rank(&self) -> usize {
140        self.shape().rank()
141    }
142
143    pub fn to_vec(&self) -> crate::Result<Vec<T>> {
144        self.iter().map(|i| i.collect())
145    }
146
147    /// Returns an iterator over **storage indices**.
148    ///
149    /// This iterator yields the linear (flat) indices as they are laid out
150    /// in the underlying storage buffer. The order depends on the memory
151    /// layout (e.g., row-major / column-major / with strides).
152    ///
153    /// Example for shape = (2, 2) in row-major layout:
154    /// yields: `0, 1, 2, 3`
155    pub fn storage_indices(&self) -> StorageIndices {
156        self.layout().storage_indices()
157    }
158
159    /// Returns an iterator over **dimension coordinates**.
160    ///
161    /// This iterator yields the multi-dimensional coordinates
162    /// (e.g., `[i, j, k, ...]`) of each element in the array, independent
163    /// of the physical storage layout.
164    ///
165    /// Example for shape = (2, 2):
166    /// yields: `[0, 0], [0, 1], [1, 0], [1, 1]`
167    pub fn dim_coordinates(&self) -> DimCoordinates {
168        self.shape().dim_coordinates()
169    }
170
171    pub fn dims_coordinates<const N: usize>(&self) -> crate::Result<DimNCoordinates<N>> {
172        self.shape().dims_coordinates::<N>()
173    }
174
175    pub fn dim2_coordinates(&self) -> crate::Result<DimNCoordinates<2>> {
176        self.shape().dim2_coordinates()
177    }
178
179    pub fn dim3_coordinates(&self) -> crate::Result<DimNCoordinates<3>> {
180        self.shape().dim3_coordinates()
181    }
182
183    pub fn dim4_coordinates(&self) -> crate::Result<DimNCoordinates<4>> {
184        self.shape().dim4_coordinates()
185    }
186
187    pub fn dim5_coordinates(&self) -> crate::Result<DimNCoordinates<5>> {
188        self.shape().dim5_coordinates()
189    }
190
191    #[inline]
192    pub fn requires_grad(&self) -> bool {
193        self.0.meta.requires_grad()
194    }
195}
196
197impl<T: NumDType> Tensor<T> {
198    pub fn allclose(&self, other: &Self, rtol: f64, atol: f64) -> crate::Result<bool> {
199        if self.shape() != other.shape() {
200            return Ok(false);
201        }
202        Ok(
203            self.iter()?.zip(other.iter()?).all(|(a, b)| a.close(b, rtol, atol))
204        )
205    }
206}
207
208impl<T: FloatDType> Tensor<T> {
209    pub fn detach(&self) -> Self {
210        if !self.requires_grad() {
211            self.clone()
212        } else {
213            Self(Arc::new(TensorImpl { 
214                id: TensorId::new(), 
215                storage: self.0.storage.clone(), 
216                layout: self.layout().clone(), 
217                meta: AutogradInfo::val(), 
218            }))
219        }
220    }
221    
222    #[inline]
223    pub fn set_requires_grad(&self, mode: bool) {
224        self.0.meta.set_requires_grad(mode);
225    }
226
227    #[inline]
228    pub fn op(&self) -> Option<&Op<T>> {
229        self.0.meta.op()
230    }
231
232    #[inline]
233    pub fn is_leaf(&self) -> bool {
234        self.0.meta.is_leaf()
235    }
236}
lumen_core/base/tensor.rs

lumen_core/base/
tensor.rs
