Skip to main content

luma_tensor/ops/construct/
float.rs

1use std::borrow::Cow;
2
3use crate::{Device, Float, FloatMeta, Result, Shape, Tensor, dtype::FloatDType};
4
5use super::helpers;
6use super::into_tensor::IntoTensor;
7use super::options::TensorCreationOptions;
8
9impl<D: Device> Tensor<D, Float> {
10    pub fn new(data: impl IntoTensor<D, Float>, device: &D) -> Result<Self> {
11        let shape = data.shape()?;
12        let storage = data.into_storage(device)?;
13        Ok(Self::from_storage(storage, shape, FloatMeta::val()))
14    }
15
16    pub fn from_vec_f64<'a, S: Into<Shape>>(data: impl Into<Cow<'a, [f64]>>, shape: S, device: &D) -> Result<Self> {
17        let shape = shape.into();
18        let storage = D::f_from_f64(data, device)?;
19        Ok(Self::from_storage(storage, shape, FloatMeta::val()))
20    }
21
22    pub fn from_vec_f32<'a, S: Into<Shape>>(data: impl Into<Cow<'a, [f32]>>, shape: S, device: &D) -> Result<Self> {
23        let shape = shape.into();
24        let storage = D::f_from_f32(data, device)?;
25        Ok(Self::from_storage(storage, shape, FloatMeta::val()))
26    }
27
28    pub fn from_slice<S: Into<Shape>>(data: &[f64], shape: S, options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
29        let options: TensorCreationOptions<D, Float> = options.into();
30        let shape = shape.into();
31        if shape.element_count() != data.len() {
32            return Err(crate::Error::ElementSizeMismatch { expected: data.len(), got: shape.element_count(), op: "from_slice" });
33        }
34        let storage = match options.dtype {
35            FloatDType::F64 => D::f_from_f64(data, &options.device)?,
36            FloatDType::F32 => {
37                let v: Vec<f32> = data.iter().map(|&x| x as f32).collect();
38                D::f_from_f32(&v, &options.device)?
39            }
40        };
41        Ok(Self::from_storage(storage, shape, FloatMeta::val()))
42    }
43
44    pub fn to_vec(&self) -> crate::Result<Vec<f64>> {
45        D::f_to_vec(&*self.storage_read()?, self.layout())
46    }
47
48    pub fn randn<S: Into<Shape>>(mean: f64, std: f64, shape: S, options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
49        let options: TensorCreationOptions<D, Float> = options.into();
50        let shape = shape.into();
51        let storage = D::f_rand_normal(&shape, mean, std, &options.device, options.dtype)?;
52        Ok(Self::from_storage(storage, shape, FloatMeta::val()))
53    }
54
55    pub fn rand<S: Into<Shape>>(lo: f64, hi: f64, shape: S, options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
56        let options: TensorCreationOptions<D, Float> = options.into();
57        let shape = shape.into();
58        let storage = D::f_rand_uniform(&shape, lo, hi, &options.device, options.dtype)?;
59        Ok(Self::from_storage(storage, shape, FloatMeta::val()))
60    }
61
62    pub fn rand_like(&self, lo: f64, hi: f64) -> Result<Self> {
63        Self::rand(lo, hi, self.shape().clone(), self.dtype())
64    }
65
66    pub fn randn_like(&self, mean: f64, std: f64) -> Result<Self> {
67        Self::randn(mean, std, self.shape().clone(), self.dtype())
68    }
69
70    pub fn eye(n: usize, options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
71        let options = options.into();
72        match options.dtype {
73            FloatDType::F64 => Self::new(helpers::fill_eye::<f64>(n), &options.device),
74            FloatDType::F32 => Self::new(helpers::fill_eye::<f32>(n), &options.device),
75        }
76    }
77
78    pub fn diag(diag: &[f64], options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
79        let options = options.into();
80        let n = diag.len();
81        match options.dtype {
82            FloatDType::F64 => {
83                let mut v = vec![0.0f64; n * n];
84                for i in 0..n {
85                    v[i * n + i] = diag[i];
86                }
87                Self::new(v, &options.device)
88            }
89            FloatDType::F32 => {
90                let mut v = vec![0.0f32; n * n];
91                for i in 0..n {
92                    v[i * n + i] = diag[i] as f32;
93                }
94                Self::new(v, &options.device)
95            }
96        }
97    }
98
99    pub fn tril(n: usize, diagonal: bool, options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
100        let options = options.into();
101        match options.dtype {
102            FloatDType::F64 => Self::new(helpers::fill_tril::<f64>(n, diagonal), &options.device),
103            FloatDType::F32 => Self::new(helpers::fill_tril::<f32>(n, diagonal), &options.device),
104        }
105    }
106
107    pub fn triu(n: usize, diagonal: bool, options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
108        let options = options.into();
109        match options.dtype {
110            FloatDType::F64 => Self::new(helpers::fill_triu::<f64>(n, diagonal), &options.device),
111            FloatDType::F32 => Self::new(helpers::fill_triu::<f32>(n, diagonal), &options.device),
112        }
113    }
114
115    pub fn linspace(start: f64, stop: f64, n: usize, options: impl Into<TensorCreationOptions<D, Float>>) -> Result<Self> {
116        let options = options.into();
117        let step = if n > 1 { (stop - start) / (n as f64 - 1.0) } else { 0.0 };
118        match options.dtype {
119            FloatDType::F64 => {
120                let v: Vec<f64> = (0..n).map(|i| start + step * i as f64).collect();
121                Self::new(v, &options.device)
122            }
123            FloatDType::F32 => {
124                let v: Vec<f32> = (0..n).map(|i| (start + step * i as f64) as f32).collect();
125                Self::new(v, &options.device)
126            }
127        }
128    }
129}