luma_tensor/ops/construct/
int.rs

1use std::borrow::Cow;
2
3use crate::{Device, Int, Result, Shape, Tensor, dtype::IntDType};
4
5use super::helpers;
6use super::into_tensor::IntoTensor;
7use super::options::TensorCreationOptions;
8
9impl<D: Device> Tensor<D, Int> {
10    pub fn new(data: impl IntoTensor<D, Int>, device: &D) -> Result<Self> {
11        let shape = data.shape()?;
12        let storage = data.into_storage(device)?;
13        Ok(Self::from_storage(storage, shape, ()))
14    }
15
16    pub fn from_vec_i64<'a, S: Into<Shape>>(data: impl Into<Cow<'a, [i64]>>, shape: S, device: &D) -> Result<Self> {
17        let shape = shape.into();
18        let storage = D::i_from_i64(data, device)?;
19        Ok(Self::from_storage(storage, shape, ()))
20    }
21
22    pub fn from_vec_i32<'a, S: Into<Shape>>(data: impl Into<Cow<'a, [i32]>>, shape: S, device: &D) -> Result<Self> {
23        let shape = shape.into();
24        let storage = D::i_from_i32(data, device)?;
25        Ok(Self::from_storage(storage, shape, ()))
26    }
27
28    pub fn from_vec_u32<'a, S: Into<Shape>>(data: impl Into<Cow<'a, [u32]>>, shape: S, device: &D) -> Result<Self> {
29        let shape = shape.into();
30        let storage = D::i_from_u32(data, device)?;
31        Ok(Self::from_storage(storage, shape, ()))
32    }
33
34    pub fn from_vec_u8<'a, S: Into<Shape>>(data: impl Into<Cow<'a, [u8]>>, shape: S, device: &D) -> Result<Self> {
35        let shape = shape.into();
36        let storage = D::i_from_u8(data, device)?;
37        Ok(Self::from_storage(storage, shape, ()))
38    }
39
40    pub fn from_slice<S: Into<Shape>>(data: &[i64], shape: S, options: impl Into<TensorCreationOptions<D, Int>>) -> Result<Self> {
41        let options: TensorCreationOptions<D, Int> = options.into();
42        let shape = shape.into();
43        if shape.element_count() != data.len() {
44            return Err(crate::Error::ElementSizeMismatch { expected: data.len(), got: shape.element_count(), op: "from_slice" });
45        }
46        let storage = match options.dtype {
47            IntDType::I32 => {
48                let v: Vec<i32> = data.iter().map(|&x| x as i32).collect();
49                D::i_from_bytes(bytemuck::cast_slice(&v), &shape, &options.device, IntDType::I32)?
50            }
51            IntDType::U32 => {
52                let v: Vec<u32> = data.iter().map(|&x| x as u32).collect();
53                D::i_from_bytes(bytemuck::cast_slice(&v), &shape, &options.device, IntDType::U32)?
54            }
55            IntDType::U8 => {
56                let v: Vec<u8> = data.iter().map(|&x| x as u8).collect();
57                D::i_from_bytes(&v, &shape, &options.device, IntDType::U8)?
58            }
59        };
60        Ok(Self::from_storage(storage, shape, ()))
61    }
62
63    pub fn to_vec(&self) -> crate::Result<Vec<i64>> {
64        D::i_to_vec(&*self.storage_read()?, self.layout())
65    }
66
67    pub fn arange(start: i64, end: i64, step: i64, options: impl Into<TensorCreationOptions<D, Int>>) -> Result<Self> {
68        let options: TensorCreationOptions<D, Int> = options.into();
69        let (storage, n) = D::i_arange(start, end, step, &options.device, options.dtype)?;
70        Ok(Self::from_storage(storage, Shape::from(n), ()))
71    }
72
73    pub fn eye(n: usize, options: impl Into<TensorCreationOptions<D, Int>>) -> Result<Self> {
74        let options = options.into();
75        match options.dtype {
76            IntDType::I32 => Self::new(helpers::fill_eye::<i32>(n), &options.device),
77            IntDType::U32 => Self::new(helpers::fill_eye::<u32>(n), &options.device),
78            IntDType::U8 => Self::new(helpers::fill_eye::<u8>(n), &options.device),
79        }
80    }
81
82    pub fn diag(diag: &[i64], options: impl Into<TensorCreationOptions<D, Int>>) -> Result<Self> {
83        let options = options.into();
84        let n = diag.len();
85        match options.dtype {
86            IntDType::I32 => {
87                let mut v = vec![0i32; n * n];
88                for i in 0..n {
89                    v[i * n + i] = diag[i] as i32;
90                }
91                Self::new(v, &options.device)
92            }
93            IntDType::U32 => {
94                let mut v = vec![0u32; n * n];
95                for i in 0..n {
96                    v[i * n + i] = diag[i] as u32;
97                }
98                Self::new(v, &options.device)
99            }
100            IntDType::U8 => {
101                let mut v = vec![0u8; n * n];
102                for i in 0..n {
103                    v[i * n + i] = diag[i] as u8;
104                }
105                Self::new(v, &options.device)
106            }
107        }
108    }
109
110    pub fn tril(n: usize, diagonal: bool, options: impl Into<TensorCreationOptions<D, Int>>) -> Result<Self> {
111        let options = options.into();
112        match options.dtype {
113            IntDType::I32 => Self::new(helpers::fill_tril::<i32>(n, diagonal), &options.device),
114            IntDType::U32 => Self::new(helpers::fill_tril::<u32>(n, diagonal), &options.device),
115            IntDType::U8 => Self::new(helpers::fill_tril::<u8>(n, diagonal), &options.device),
116        }
117    }
118
119    pub fn triu(n: usize, diagonal: bool, options: impl Into<TensorCreationOptions<D, Int>>) -> Result<Self> {
120        let options = options.into();
121        match options.dtype {
122            IntDType::I32 => Self::new(helpers::fill_triu::<i32>(n, diagonal), &options.device),
123            IntDType::U32 => Self::new(helpers::fill_triu::<u32>(n, diagonal), &options.device),
124            IntDType::U8 => Self::new(helpers::fill_triu::<u8>(n, diagonal), &options.device),
125        }
126    }
127}
luma_tensor/ops/construct/int.rs

luma_tensor/ops/construct/
int.rs
