Skip to main content

luma_tensor/ops/construct/
from_element.rs

1use std::borrow::Cow;
2
3use crate::{Device, Result};
4
5pub trait MatrixFill: Clone + Sized {
6    fn zero_val() -> Self;
7    fn one_val() -> Self;
8}
9
10impl MatrixFill for f64 {
11    fn zero_val() -> Self {
12        0.0
13    }
14    fn one_val() -> Self {
15        1.0
16    }
17}
18impl MatrixFill for f32 {
19    fn zero_val() -> Self {
20        0.0_f32
21    }
22    fn one_val() -> Self {
23        1.0_f32
24    }
25}
26impl MatrixFill for i64 {
27    fn zero_val() -> Self {
28        0
29    }
30    fn one_val() -> Self {
31        1
32    }
33}
34impl MatrixFill for i32 {
35    fn zero_val() -> Self {
36        0
37    }
38    fn one_val() -> Self {
39        1
40    }
41}
42impl MatrixFill for u32 {
43    fn zero_val() -> Self {
44        0
45    }
46    fn one_val() -> Self {
47        1
48    }
49}
50impl MatrixFill for u8 {
51    fn zero_val() -> Self {
52        0
53    }
54    fn one_val() -> Self {
55        1
56    }
57}
58impl MatrixFill for bool {
59    fn zero_val() -> Self {
60        false
61    }
62    fn one_val() -> Self {
63        true
64    }
65}
66
67pub trait FloatFrom: bytemuck::Pod + Copy + MatrixFill + 'static {
68    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::FloatStorage>;
69}
70
71impl FloatFrom for f64 {
72    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::FloatStorage> {
73        D::f_from_f64(data, device)
74    }
75}
76
77impl FloatFrom for f32 {
78    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::FloatStorage> {
79        D::f_from_f32(data, device)
80    }
81}
82
83pub trait IntFrom: bytemuck::Pod + Copy + MatrixFill + 'static {
84    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::IntStorage>;
85}
86
87impl IntFrom for i64 {
88    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::IntStorage> {
89        D::i_from_i64(data, device)
90    }
91}
92
93impl IntFrom for i32 {
94    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::IntStorage> {
95        D::i_from_i32(data, device)
96    }
97}
98
99impl IntFrom for u32 {
100    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::IntStorage> {
101        D::i_from_u32(data, device)
102    }
103}
104
105impl IntFrom for u8 {
106    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::IntStorage> {
107        D::i_from_u8(data, device)
108    }
109}
110
111pub trait BoolFrom: Copy + MatrixFill + 'static {
112    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::BoolStorage>;
113}
114
115impl BoolFrom for bool {
116    fn into_storage<'a, D: Device>(data: impl Into<Cow<'a, [Self]>>, device: &D) -> Result<D::BoolStorage> {
117        D::b_from_bool(data, device)
118    }
119}