luma_tensor/ops/construct/
float.rs1use 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}