burn_std/element/
base.rs

1use core::cmp::Ordering;
2use rand::Rng;
3
4use crate::distribution::Distribution;
5use crate::{BoolStore, DType, bf16, f16, flex32};
6
7use super::cast::ToElement;
8
9/// Core element trait for tensor values.
10///
11/// This trait defines the minimal set of capabilities required for a type to be
12/// stored and manipulated as a tensor element across all backends.
13pub trait Element:
14    ToElement
15    + ElementRandom
16    + ElementConversion
17    + ElementEq
18    + bytemuck::CheckedBitPattern
19    + bytemuck::NoUninit
20    + bytemuck::Zeroable
21    + core::fmt::Debug
22    + core::fmt::Display
23    + Default
24    + Send
25    + Sync
26    + Copy
27    + 'static
28{
29    /// The dtype of the element.
30    fn dtype() -> DType;
31}
32
33/// Ordered element trait for tensor values.
34///
35/// This trait extends [`Element`] with ordering semantics, enabling comparison
36/// and order-dependent operations in generic Rust implementations.
37///
38/// Backends that implement these operations entirely at the device level do
39/// not rely on this trait. It only constrains the scalar type for generic Rust code.
40pub trait ElementOrdered: Element + ElementComparison + ElementLimits {}
41
42/// Element conversion trait for tensor.
43pub trait ElementConversion {
44    /// Converts an element to another element.
45    ///
46    /// # Arguments
47    ///
48    /// * `elem` - The element to convert.
49    ///
50    /// # Returns
51    ///
52    /// The converted element.
53    fn from_elem<E: ToElement>(elem: E) -> Self;
54
55    /// Converts and returns the converted element.
56    fn elem<E: Element>(self) -> E;
57}
58
59/// Element trait for random value of a tensor.
60pub trait ElementRandom {
61    /// Returns a random value for the given distribution.
62    ///
63    /// # Arguments
64    ///
65    /// * `distribution` - The distribution to sample from.
66    /// * `rng` - The random number generator.
67    ///
68    /// # Returns
69    ///
70    /// The random value.
71    fn random<R: Rng>(distribution: Distribution, rng: &mut R) -> Self;
72}
73
74/// Element trait for equality of a tensor.
75pub trait ElementEq {
76    /// Returns whether `self` and `other` are equal.
77    fn eq(&self, other: &Self) -> bool;
78}
79
80/// Element ordering trait.
81pub trait ElementComparison {
82    /// Returns and [Ordering] between `self` and `other`.
83    fn cmp(&self, other: &Self) -> Ordering;
84}
85
86/// Element limits trait.
87pub trait ElementLimits {
88    /// The minimum representable value
89    const MIN: Self;
90    /// The maximum representable value
91    const MAX: Self;
92}
93
94/// Macro to implement the element trait for a type.
95#[macro_export]
96macro_rules! make_element {
97    (
98        ty $type:ident,
99        convert $convert:expr,
100        random $random:expr,
101        cmp $cmp:expr,
102        dtype $dtype:expr
103    ) => {
104        make_element!(ty $type, convert $convert, random $random, cmp $cmp, dtype $dtype, min $type::MIN, max $type::MAX);
105    };
106    (
107        ty $type:ident,
108        convert $convert:expr,
109        random $random:expr,
110        cmp $cmp:expr,
111        dtype $dtype:expr,
112        min $min:expr,
113        max $max:expr
114    ) => {
115        impl Element for $type {
116            #[inline(always)]
117            fn dtype() -> $crate::DType {
118                $dtype
119            }
120        }
121        impl ElementEq for $type {
122            fn eq(&self, other: &Self) -> bool {
123                self == other
124            }
125        }
126
127        impl ElementConversion for $type {
128            #[inline(always)]
129            fn from_elem<E: ToElement>(elem: E) -> Self {
130                #[allow(clippy::redundant_closure_call)]
131                $convert(&elem)
132            }
133            #[inline(always)]
134            fn elem<E: Element>(self) -> E {
135                E::from_elem(self)
136            }
137        }
138
139        impl ElementRandom for $type {
140            fn random<R: Rng>(distribution: Distribution, rng: &mut R) -> Self {
141                #[allow(clippy::redundant_closure_call)]
142                $random(distribution, rng)
143            }
144        }
145
146        impl ElementComparison for $type {
147            fn cmp(&self, other: &Self) -> Ordering {
148                let a = self.elem::<$type>();
149                let b = other.elem::<$type>();
150                #[allow(clippy::redundant_closure_call)]
151                $cmp(&a, &b)
152            }
153        }
154
155        impl ElementLimits for $type {
156            const MIN: Self = $min;
157            const MAX: Self = $max;
158        }
159
160        impl ElementOrdered for $type {}
161
162    };
163}
164
165make_element!(
166    ty f64,
167    convert ToElement::to_f64,
168    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
169    cmp |a: &f64, b: &f64| a.total_cmp(b),
170    dtype DType::F64
171);
172
173make_element!(
174    ty f32,
175    convert ToElement::to_f32,
176    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
177    cmp |a: &f32, b: &f32| a.total_cmp(b),
178    dtype DType::F32
179);
180
181make_element!(
182    ty i64,
183    convert ToElement::to_i64,
184    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
185    cmp |a: &i64, b: &i64| Ord::cmp(a, b),
186    dtype DType::I64
187);
188
189make_element!(
190    ty u64,
191    convert ToElement::to_u64,
192    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
193    cmp |a: &u64, b: &u64| Ord::cmp(a, b),
194    dtype DType::U64
195);
196
197make_element!(
198    ty i32,
199    convert ToElement::to_i32,
200    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
201    cmp |a: &i32, b: &i32| Ord::cmp(a, b),
202    dtype DType::I32
203);
204
205make_element!(
206    ty u32,
207    convert ToElement::to_u32,
208    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
209    cmp |a: &u32, b: &u32| Ord::cmp(a, b),
210    dtype DType::U32
211);
212
213make_element!(
214    ty i16,
215    convert ToElement::to_i16,
216    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
217    cmp |a: &i16, b: &i16| Ord::cmp(a, b),
218    dtype DType::I16
219);
220
221make_element!(
222    ty u16,
223    convert ToElement::to_u16,
224    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
225    cmp |a: &u16, b: &u16| Ord::cmp(a, b),
226    dtype DType::U16
227);
228
229make_element!(
230    ty i8,
231    convert ToElement::to_i8,
232    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
233    cmp |a: &i8, b: &i8| Ord::cmp(a, b),
234    dtype DType::I8
235);
236
237make_element!(
238    ty u8,
239    convert ToElement::to_u8,
240    random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(),
241    cmp |a: &u8, b: &u8| Ord::cmp(a, b),
242    dtype DType::U8
243);
244
245make_element!(
246    ty f16,
247    convert ToElement::to_f16,
248    random |distribution: Distribution, rng: &mut R| {
249        let sample: f32 = distribution.sampler(rng).sample();
250        f16::from_elem(sample)
251    },
252    cmp |a: &f16, b: &f16| a.total_cmp(b),
253    dtype DType::F16
254);
255make_element!(
256    ty bf16,
257    convert ToElement::to_bf16,
258    random |distribution: Distribution, rng: &mut R| {
259        let sample: f32 = distribution.sampler(rng).sample();
260        bf16::from_elem(sample)
261    },
262    cmp |a: &bf16, b: &bf16| a.total_cmp(b),
263    dtype DType::BF16
264);
265
266make_element!(
267    ty flex32,
268    convert |elem: &dyn ToElement| flex32::from_f32(elem.to_f32()),
269    random |distribution: Distribution, rng: &mut R| {
270        let sample: f32 = distribution.sampler(rng).sample();
271        flex32::from_elem(sample)
272    },
273    cmp |a: &flex32, b: &flex32| a.total_cmp(b),
274    dtype DType::Flex32,
275    min flex32::from_f32(f16::MIN.to_f32_const()),
276    max flex32::from_f32(f16::MAX.to_f32_const())
277);
278
279make_element!(
280    ty bool,
281    convert ToElement::to_bool,
282    random |distribution: Distribution, rng: &mut R| {
283        let sample: u8 = distribution.sampler(rng).sample();
284        bool::from_elem(sample)
285    },
286    cmp |a: &bool, b: &bool| Ord::cmp(a, b),
287    dtype DType::Bool(BoolStore::Native),
288    min false,
289    max true
290);
burn_std/element/base.rs

burn_std/element/
base.rs
