lumen_core/ops/
convert.rs

1use crate::{AutogradMetaT, DTypeConvert, Error, Result, TensorOrScalar, WithDType};
2use crate::Tensor;
3
4impl<T: WithDType> Tensor<T> {
5    pub fn contiguous(&self) -> crate::Result<Tensor<T>> {
6        if self.is_contiguous() {
7            Ok(self.clone())
8        } else {
9            self.copy()
10        }
11    }
12
13    pub fn copy(&self) -> crate::Result<Self> {
14        // MARK? need meta？
15        let storage = self.storage_read()?.copy(self.layout());
16        let meta = T::AutogradMeta::on_copy_op(self);
17        Ok(Self::from_storage(storage, self.shape(), meta))
18    }
19
20    /// WARNING: This is an unsafe operation that bypasses the autograd engine. 
21    /// It does NOT track computation history and cannot compute gradients. 
22    /// Attempting to call this method on a tensor with `requires_grad=True` 
23    /// will raise a Error.
24    pub fn copy_(&self, source: impl Into<TensorOrScalar<T>>) -> Result<()> {
25        self.check_implace_op()?;
26        self.impl_copy_(source)
27    }
28
29    pub(crate) fn impl_copy_(&self, source: impl Into<TensorOrScalar<T>>) -> Result<()> {
30        match source.into() {
31            TensorOrScalar::Scalar(src) => {
32                let mut storage = self.storage_write()?;
33                for storage_index in self.layout().storage_indices() {
34                    storage.set_unchecked(storage_index, src);
35                }
36                Ok(())
37            }
38            TensorOrScalar::Tensor(src) => {
39                if src.shape() != self.shape() {
40                    Err(Error::ShapeMismatchCopyFrom { dst: self.shape().clone(), src: src.shape().clone() })?
41                }
42        
43                let mut storage = self.storage_write()?;
44                for (self_storage_index, src_value) in self.layout().storage_indices().zip(src.iter()?) {
45                    storage.set_unchecked(self_storage_index, src_value);
46                }
47        
48                Ok(())
49            }
50        }
51    }
52}
53
54impl<From: WithDType> Tensor<From> {
55    pub fn cast<To: WithDType>(&self) -> crate::Result<Tensor<To>> 
56    where
57        From: DTypeConvert<To>,
58    {
59        // if TypeId::of::<From>() == TypeId::of::<To>() {
60        //     let self_any = self as &dyn Any;            
61        //     if let Some(same_tensor) = self_any.downcast_ref::<Tensor<To>>() {
62        //         return same_tensor.clone();
63        //     }
64        // }
65        let storage = self.storage_read()?.copy_map(self.layout(), From::convert);
66        Ok(Tensor::<To>::from_storage(storage, self.shape(), Default::default()))
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use crate::IndexOp;
73
74    use super::*;
75
76    #[test]
77    fn test_copy_() {
78        let a = Tensor::new(&[[1, 2, 3], [3, 4, 5], [4, 5, 6]]).unwrap();
79        a.index(0).unwrap().copy_(100).unwrap();
80        println!("{}", a);
81        a.index((1, 1)).unwrap().copy_(200).unwrap();
82        println!("{}", a);
83        a.index((1.., 1..)).unwrap().copy_(999).unwrap();
84        println!("{}", a);
85    }
86
87    #[test]
88    fn test_copy_1d() {
89        let a = Tensor::new(&[1, 2, 3]).unwrap();
90        let b = a.copy().unwrap();
91        assert_eq!(a.shape(), b.shape());
92        assert_eq!(a.to_vec().unwrap(), b.to_vec().unwrap());
93    }
94
95    #[test]
96    fn test_copy_2d() {
97        let a = Tensor::new(&[[1, 2], [3, 4]]).unwrap();
98        let b = a.copy().unwrap();
99        assert_eq!(a.shape(), b.shape());
100        assert_eq!(a.to_vec().unwrap(), b.to_vec().unwrap());
101    }
102
103    #[test]
104    fn test_cast_i32_to_f32() {
105        let a = Tensor::new(&[1i32, 2, 3]).unwrap();
106        let b: Tensor<f32> = a.cast().unwrap();
107        assert_eq!(b.shape(), a.shape());
108        let expected = Tensor::new(&[1.0f32, 2.0, 3.0]).unwrap();
109        assert!(b.allclose(&expected, 1e-6, 1e-6).unwrap());
110    }
111
112    #[test]
113    fn test_cast_f64_to_i32() {
114        let a = Tensor::new(&[1.5f64, 2.7, 3.0]).unwrap();
115        let b: Tensor<i32> = a.cast().unwrap();
116        assert_eq!(b.shape(), a.shape());
117        let expected = Tensor::new(&[1i32, 2, 3]).unwrap(); // cast truncates
118        assert_eq!(b.to_vec().unwrap(), expected.to_vec().unwrap());
119    }
120
121    #[test]
122    fn test_cast_2d() {
123        let a = Tensor::new(&[[1i32, 2], [3, 4]]).unwrap();
124        let b: Tensor<f64> = a.cast().unwrap();
125        assert_eq!(b.shape(), a.shape());
126        let expected = Tensor::new(&[[1.0, 2.0], [3.0, 4.0]]).unwrap();
127        assert!(b.allclose(&expected, 1e-12, 1e-12).unwrap());
128    }
129
130    #[test]
131    fn test_copy_vs_cast() {
132        let a = Tensor::new(&[1i32, 2, 3]).unwrap();
133        let b = a.copy().unwrap();
134        let c: Tensor<f32> = a.cast().unwrap();
135
136        assert_eq!(b.to_vec().unwrap(), a.to_vec().unwrap());
137
138        let expected = Tensor::new(&[1.0f32, 2.0, 3.0]).unwrap();
139        assert!(c.allclose(&expected, 1e-6, 1e-6).unwrap());
140    }
141}
lumen_core/ops/convert.rs

lumen_core/ops/
convert.rs
