fuel_tensor/
dtype.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Types for elements that can be stored and manipulated using tensors.
3//!
4//! The [`DType`] enum and its inherent methods, Display, FromStr, DTypeParseError,
5//! and safetensors interop are all defined in `fuel-core-types` and re-exported here.
6//!
7//! [`WithDType`] is a local marker subtrait of `fuel_ir::dtype::WithDType`.
8//! Keeping it local lets Rust's coherence checker prove disjointness of generic impls
9//! (e.g. `NdArray for S` vs `NdArray for Vec<S>`).  All methods are inherited from
10//! the upstream trait.
11#![allow(clippy::redundant_closure_call)]
12
13// Re-export DType and related types from fuel-core-types.
14pub use fuel_ir::dtype::{DType, DTypeParseError};
15
16/// Local marker subtrait — inherits every method from
17/// `fuel_ir::dtype::WithDType` but is defined in this crate so that
18/// the orphan/coherence rules treat it as a local trait.
19///
20/// Also requires `fuel_ir::HostDType` (the host-buffer conversion methods split
21/// off `WithDType` in B0.4): every `BuiltinDType` here is host-storable, so this
22/// keeps `crate::WithDType`-bounded fuel-core code (NdArray, tensor `to_vec*` /
23/// `from_*`) able to call `to_cpu_storage` / `cpu_storage_as_slice` unchanged.
24pub trait WithDType: fuel_ir::dtype::WithDType + fuel_ir::HostDType {}
25
26macro_rules! with_dtype {
27    ($ty:ty, $dtype:ident, $from_f64:expr, $to_f64:expr) => {
28        impl WithDType for $ty {}
29    };
30}
31use float8::F8E4M3 as f8e4m3;
32use half::{bf16, f16};
33
34with_dtype!(u8, U8, |v: f64| v as u8, |v: u8| v as f64);
35with_dtype!(i8, I8, |v: f64| v as i8, |v: i8| v as f64);
36with_dtype!(u32, U32, |v: f64| v as u32, |v: u32| v as f64);
37with_dtype!(i16, I16, |v: f64| v as i16, |v: i16| v as f64);
38with_dtype!(i32, I32, |v: f64| v as i32, |v: i32| v as f64);
39with_dtype!(i64, I64, |v: f64| v as i64, |v: i64| v as f64);
40with_dtype!(f16, F16, f16::from_f64, f16::to_f64);
41with_dtype!(bf16, BF16, bf16::from_f64, bf16::to_f64);
42with_dtype!(f32, F32, |v: f64| v as f32, |v: f32| v as f64);
43with_dtype!(f64, F64, |v: f64| v, |v: f64| v);
44with_dtype!(f8e4m3, F8E4M3, f8e4m3::from_f64, |v: f8e4m3| v.to_f64());
45
46pub trait IntDType: WithDType + num_traits::Bounded {
47    fn is_true(&self) -> bool;
48    fn as_usize(&self) -> usize;
49}
50
51impl IntDType for i64 {
52    fn is_true(&self) -> bool {
53        *self != 0
54    }
55    fn as_usize(&self) -> usize {
56        *self as usize
57    }
58}
59
60impl IntDType for u32 {
61    fn is_true(&self) -> bool {
62        *self != 0
63    }
64    fn as_usize(&self) -> usize {
65        *self as usize
66    }
67}
68
69impl IntDType for u8 {
70    fn is_true(&self) -> bool {
71        *self != 0
72    }
73    fn as_usize(&self) -> usize {
74        *self as usize
75    }
76}
77
78impl IntDType for i8 {
79    fn is_true(&self) -> bool {
80        *self != 0
81    }
82    fn as_usize(&self) -> usize {
83        *self as usize
84    }
85}
86
87impl IntDType for i16 {
88    fn is_true(&self) -> bool {
89        *self != 0
90    }
91    fn as_usize(&self) -> usize {
92        *self as usize
93    }
94}
95
96impl IntDType for i32 {
97    fn is_true(&self) -> bool {
98        *self != 0
99    }
100    fn as_usize(&self) -> usize {
101        *self as usize
102    }
103}
104
105pub trait FloatDType: WithDType {}
106
107impl FloatDType for f16 {}
108impl FloatDType for bf16 {}
109impl FloatDType for f32 {}
110impl FloatDType for f64 {}
111impl FloatDType for f8e4m3 {}
fuel_tensor/dtype.rs

fuel_tensor/
dtype.rs
