Skip to main content

luma_tensor/ops/construct/
into_tensor.rs

1use super::from_element::{BoolFrom, FloatFrom, IntFrom};
2use crate::{Bool, DTypeKind, Device, Float, Int, Result, Shape};
3
4pub trait IntoTensor<D: Device, K: DTypeKind<D>> {
5    fn shape(&self) -> Result<Shape>;
6    fn into_storage(self, device: &D) -> Result<K::Storage>;
7}
8
9// ---- Float impls --------------------------------------------------------------
10
11impl<D: Device, T: FloatFrom> IntoTensor<D, Float> for T {
12    fn shape(&self) -> Result<Shape> {
13        Ok(Shape::scalar())
14    }
15    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
16        T::into_storage(&[self][..], device)
17    }
18}
19
20impl<D: Device, T: FloatFrom> IntoTensor<D, Float> for &[T] {
21    fn shape(&self) -> Result<Shape> {
22        Ok(Shape::from(self.len()))
23    }
24    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
25        T::into_storage(self, device)
26    }
27}
28
29impl<D: Device, T: FloatFrom> IntoTensor<D, Float> for Vec<T> {
30    fn shape(&self) -> Result<Shape> {
31        Ok(Shape::from(self.len()))
32    }
33    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
34        T::into_storage(self, device)
35    }
36}
37
38impl<D: Device, T: FloatFrom, const N: usize> IntoTensor<D, Float> for [T; N] {
39    fn shape(&self) -> Result<Shape> {
40        Ok(Shape::from(N))
41    }
42    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
43        T::into_storage(&self[..], device)
44    }
45}
46
47impl<D: Device, T: FloatFrom, const N: usize> IntoTensor<D, Float> for &[T; N] {
48    fn shape(&self) -> Result<Shape> {
49        Ok(Shape::from(N))
50    }
51    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
52        T::into_storage(self.as_slice(), device)
53    }
54}
55
56impl<D: Device, T: FloatFrom, const R: usize, const C: usize> IntoTensor<D, Float> for &[[T; C]; R] {
57    fn shape(&self) -> Result<Shape> {
58        Ok(Shape::from((R, C)))
59    }
60    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
61        let v = self.concat();
62        T::into_storage(&v, device)
63    }
64}
65
66impl<D: Device, T: FloatFrom, const D1: usize, const D2: usize, const D3: usize> IntoTensor<D, Float> for &[[[T; D3]; D2]; D1] {
67    fn shape(&self) -> Result<Shape> {
68        Ok(Shape::from((D1, D2, D3)))
69    }
70    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
71        let mut v = Vec::with_capacity(D1 * D2 * D3);
72        for i1 in 0..D1 {
73            for i2 in 0..D2 {
74                v.extend_from_slice(&self[i1][i2]);
75            }
76        }
77        T::into_storage(v, device)
78    }
79}
80
81impl<D: Device, T: FloatFrom, const D1: usize, const D2: usize, const D3: usize, const D4: usize> IntoTensor<D, Float>
82    for &[[[[T; D4]; D3]; D2]; D1]
83{
84    fn shape(&self) -> Result<Shape> {
85        Ok(Shape::from((D1, D2, D3, D4)))
86    }
87    fn into_storage(self, device: &D) -> Result<D::FloatStorage> {
88        let mut v = Vec::with_capacity(D1 * D2 * D3 * D4);
89        for i1 in 0..D1 {
90            for i2 in 0..D2 {
91                for i3 in 0..D3 {
92                    v.extend_from_slice(&self[i1][i2][i3]);
93                }
94            }
95        }
96        T::into_storage(v, device)
97    }
98}
99
100// ---- Int impls ----------------------------------------------------------------
101
102impl<D: Device, T: IntFrom> IntoTensor<D, Int> for T {
103    fn shape(&self) -> Result<Shape> {
104        Ok(Shape::scalar())
105    }
106    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
107        T::into_storage(&[self][..], device)
108    }
109}
110
111impl<D: Device, T: IntFrom> IntoTensor<D, Int> for &[T] {
112    fn shape(&self) -> Result<Shape> {
113        Ok(Shape::from(self.len()))
114    }
115    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
116        T::into_storage(self, device)
117    }
118}
119
120impl<D: Device, T: IntFrom> IntoTensor<D, Int> for Vec<T> {
121    fn shape(&self) -> Result<Shape> {
122        Ok(Shape::from(self.len()))
123    }
124    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
125        T::into_storage(self, device)
126    }
127}
128
129impl<D: Device, T: IntFrom, const N: usize> IntoTensor<D, Int> for [T; N] {
130    fn shape(&self) -> Result<Shape> {
131        Ok(Shape::from(N))
132    }
133    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
134        T::into_storage(&self[..], device)
135    }
136}
137
138impl<D: Device, T: IntFrom, const N: usize> IntoTensor<D, Int> for &[T; N] {
139    fn shape(&self) -> Result<Shape> {
140        Ok(Shape::from(N))
141    }
142    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
143        T::into_storage(self.as_slice(), device)
144    }
145}
146
147impl<D: Device, T: IntFrom, const R: usize, const C: usize> IntoTensor<D, Int> for &[[T; C]; R] {
148    fn shape(&self) -> Result<Shape> {
149        Ok(Shape::from((R, C)))
150    }
151    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
152        let v = self.concat();
153        T::into_storage(&v, device)
154    }
155}
156
157impl<D: Device, T: IntFrom, const D1: usize, const D2: usize, const D3: usize> IntoTensor<D, Int> for &[[[T; D3]; D2]; D1] {
158    fn shape(&self) -> Result<Shape> {
159        Ok(Shape::from((D1, D2, D3)))
160    }
161    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
162        let mut v = Vec::with_capacity(D1 * D2 * D3);
163        for i1 in 0..D1 {
164            for i2 in 0..D2 {
165                v.extend_from_slice(&self[i1][i2]);
166            }
167        }
168        T::into_storage(v, device)
169    }
170}
171
172impl<D: Device, T: IntFrom, const D1: usize, const D2: usize, const D3: usize, const D4: usize> IntoTensor<D, Int>
173    for &[[[[T; D4]; D3]; D2]; D1]
174{
175    fn shape(&self) -> Result<Shape> {
176        Ok(Shape::from((D1, D2, D3, D4)))
177    }
178    fn into_storage(self, device: &D) -> Result<D::IntStorage> {
179        let mut v = Vec::with_capacity(D1 * D2 * D3 * D4);
180        for i1 in 0..D1 {
181            for i2 in 0..D2 {
182                for i3 in 0..D3 {
183                    v.extend_from_slice(&self[i1][i2][i3]);
184                }
185            }
186        }
187        T::into_storage(v, device)
188    }
189}
190
191// ---- Bool impls ---------------------------------------------------------------
192
193impl<D: Device, T: BoolFrom> IntoTensor<D, Bool> for T {
194    fn shape(&self) -> Result<Shape> {
195        Ok(Shape::scalar())
196    }
197    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
198        T::into_storage(&[self][..], device)
199    }
200}
201
202impl<D: Device, T: BoolFrom> IntoTensor<D, Bool> for &[T] {
203    fn shape(&self) -> Result<Shape> {
204        Ok(Shape::from(self.len()))
205    }
206    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
207        T::into_storage(self, device)
208    }
209}
210
211impl<D: Device, T: BoolFrom> IntoTensor<D, Bool> for Vec<T> {
212    fn shape(&self) -> Result<Shape> {
213        Ok(Shape::from(self.len()))
214    }
215    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
216        T::into_storage(self, device)
217    }
218}
219
220impl<D: Device, T: BoolFrom, const N: usize> IntoTensor<D, Bool> for [T; N] {
221    fn shape(&self) -> Result<Shape> {
222        Ok(Shape::from(N))
223    }
224    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
225        T::into_storage(&self[..], device)
226    }
227}
228
229impl<D: Device, T: BoolFrom, const N: usize> IntoTensor<D, Bool> for &[T; N] {
230    fn shape(&self) -> Result<Shape> {
231        Ok(Shape::from(N))
232    }
233    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
234        T::into_storage(self.as_slice(), device)
235    }
236}
237
238impl<D: Device, T: BoolFrom, const R: usize, const C: usize> IntoTensor<D, Bool> for &[[T; C]; R] {
239    fn shape(&self) -> Result<Shape> {
240        Ok(Shape::from((R, C)))
241    }
242    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
243        let v = self.concat();
244        T::into_storage(&v, device)
245    }
246}
247
248impl<D: Device, T: BoolFrom, const D1: usize, const D2: usize, const D3: usize> IntoTensor<D, Bool> for &[[[T; D3]; D2]; D1] {
249    fn shape(&self) -> Result<Shape> {
250        Ok(Shape::from((D1, D2, D3)))
251    }
252    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
253        let mut v = Vec::with_capacity(D1 * D2 * D3);
254        for i1 in 0..D1 {
255            for i2 in 0..D2 {
256                v.extend_from_slice(&self[i1][i2]);
257            }
258        }
259        T::into_storage(v, device)
260    }
261}
262
263impl<D: Device, T: BoolFrom, const D1: usize, const D2: usize, const D3: usize, const D4: usize> IntoTensor<D, Bool>
264    for &[[[[T; D4]; D3]; D2]; D1]
265{
266    fn shape(&self) -> Result<Shape> {
267        Ok(Shape::from((D1, D2, D3, D4)))
268    }
269    fn into_storage(self, device: &D) -> Result<D::BoolStorage> {
270        let mut v = Vec::with_capacity(D1 * D2 * D3 * D4);
271        for i1 in 0..D1 {
272            for i2 in 0..D2 {
273                for i3 in 0..D3 {
274                    v.extend_from_slice(&self[i1][i2][i3]);
275                }
276            }
277        }
278        T::into_storage(v, device)
279    }
280}