550 lines
10 KiB
Rust
550 lines
10 KiB
Rust
//! Element type traits for tensor data.
|
|
//!
|
|
//! These traits define the numeric types that can be stored in tensors.
|
|
//! Each backend specifies which element types it supports.
|
|
|
|
use half::{bf16, f16};
|
|
// num-traits import for potential future use
|
|
#[allow(unused_imports)]
|
|
use num_traits::ToPrimitive;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::fmt::{Debug, Display};
|
|
|
|
/// Trait for floating-point element types.
|
|
///
|
|
/// Supported types:
|
|
/// - `f32` - Standard single precision
|
|
/// - `f64` - Double precision
|
|
/// - `f16` - Half precision (IEEE 754)
|
|
/// - `bf16` - Brain floating point (Google format)
|
|
///
|
|
/// # Performance Notes
|
|
///
|
|
/// - Use `f32` as the default for training
|
|
/// - Use `f16`/`bf16` for inference to reduce memory
|
|
/// - `bf16` has better dynamic range than `f16` but lower precision
|
|
pub trait FloatElement:
|
|
Clone
|
|
+ Copy
|
|
+ Debug
|
|
+ Default
|
|
+ Display
|
|
+ PartialEq
|
|
+ PartialOrd
|
|
+ Send
|
|
+ Sync
|
|
+ Serialize
|
|
+ for<'de> Deserialize<'de>
|
|
+ 'static
|
|
{
|
|
/// The name of this element type.
|
|
fn name() -> &'static str;
|
|
|
|
/// Number of bytes per element.
|
|
fn bytes() -> usize;
|
|
|
|
/// Create from f64.
|
|
fn from_f64(val: f64) -> Self;
|
|
|
|
/// Convert to f64.
|
|
fn to_f64(self) -> f64;
|
|
|
|
/// Create from f32.
|
|
fn from_f32(val: f32) -> Self;
|
|
|
|
/// Convert to f32.
|
|
fn to_f32(self) -> f32;
|
|
|
|
/// Zero value.
|
|
fn zero() -> Self;
|
|
|
|
/// One value.
|
|
fn one() -> Self;
|
|
|
|
/// Minimum representable value.
|
|
fn min_value() -> Self;
|
|
|
|
/// Maximum representable value.
|
|
fn max_value() -> Self;
|
|
|
|
/// Not a number.
|
|
fn nan() -> Self;
|
|
|
|
/// Positive infinity.
|
|
fn infinity() -> Self;
|
|
|
|
/// Negative infinity.
|
|
fn neg_infinity() -> Self;
|
|
|
|
/// Check if NaN.
|
|
fn is_nan(self) -> bool;
|
|
|
|
/// Check if infinite.
|
|
fn is_infinite(self) -> bool;
|
|
|
|
/// Epsilon value.
|
|
fn epsilon() -> Self;
|
|
}
|
|
|
|
impl FloatElement for f32 {
|
|
#[inline]
|
|
fn name() -> &'static str {
|
|
"f32"
|
|
}
|
|
#[inline]
|
|
fn bytes() -> usize {
|
|
4
|
|
}
|
|
#[inline]
|
|
fn from_f64(val: f64) -> Self {
|
|
val as f32
|
|
}
|
|
#[inline]
|
|
fn to_f64(self) -> f64 {
|
|
self as f64
|
|
}
|
|
#[inline]
|
|
fn from_f32(val: f32) -> Self {
|
|
val
|
|
}
|
|
#[inline]
|
|
fn to_f32(self) -> f32 {
|
|
self
|
|
}
|
|
#[inline]
|
|
fn zero() -> Self {
|
|
0.0
|
|
}
|
|
#[inline]
|
|
fn one() -> Self {
|
|
1.0
|
|
}
|
|
#[inline]
|
|
fn min_value() -> Self {
|
|
f32::MIN
|
|
}
|
|
#[inline]
|
|
fn max_value() -> Self {
|
|
f32::MAX
|
|
}
|
|
#[inline]
|
|
fn nan() -> Self {
|
|
f32::NAN
|
|
}
|
|
#[inline]
|
|
fn infinity() -> Self {
|
|
f32::INFINITY
|
|
}
|
|
#[inline]
|
|
fn neg_infinity() -> Self {
|
|
f32::NEG_INFINITY
|
|
}
|
|
#[inline]
|
|
fn is_nan(self) -> bool {
|
|
f32::is_nan(self)
|
|
}
|
|
#[inline]
|
|
fn is_infinite(self) -> bool {
|
|
f32::is_infinite(self)
|
|
}
|
|
#[inline]
|
|
fn epsilon() -> Self {
|
|
f32::EPSILON
|
|
}
|
|
}
|
|
|
|
impl FloatElement for f64 {
|
|
#[inline]
|
|
fn name() -> &'static str {
|
|
"f64"
|
|
}
|
|
#[inline]
|
|
fn bytes() -> usize {
|
|
8
|
|
}
|
|
#[inline]
|
|
fn from_f64(val: f64) -> Self {
|
|
val
|
|
}
|
|
#[inline]
|
|
fn to_f64(self) -> f64 {
|
|
self
|
|
}
|
|
#[inline]
|
|
fn from_f32(val: f32) -> Self {
|
|
val as f64
|
|
}
|
|
#[inline]
|
|
fn to_f32(self) -> f32 {
|
|
self as f32
|
|
}
|
|
#[inline]
|
|
fn zero() -> Self {
|
|
0.0
|
|
}
|
|
#[inline]
|
|
fn one() -> Self {
|
|
1.0
|
|
}
|
|
#[inline]
|
|
fn min_value() -> Self {
|
|
f64::MIN
|
|
}
|
|
#[inline]
|
|
fn max_value() -> Self {
|
|
f64::MAX
|
|
}
|
|
#[inline]
|
|
fn nan() -> Self {
|
|
f64::NAN
|
|
}
|
|
#[inline]
|
|
fn infinity() -> Self {
|
|
f64::INFINITY
|
|
}
|
|
#[inline]
|
|
fn neg_infinity() -> Self {
|
|
f64::NEG_INFINITY
|
|
}
|
|
#[inline]
|
|
fn is_nan(self) -> bool {
|
|
f64::is_nan(self)
|
|
}
|
|
#[inline]
|
|
fn is_infinite(self) -> bool {
|
|
f64::is_infinite(self)
|
|
}
|
|
#[inline]
|
|
fn epsilon() -> Self {
|
|
f64::EPSILON
|
|
}
|
|
}
|
|
|
|
impl FloatElement for f16 {
|
|
#[inline]
|
|
fn name() -> &'static str {
|
|
"f16"
|
|
}
|
|
#[inline]
|
|
fn bytes() -> usize {
|
|
2
|
|
}
|
|
#[inline]
|
|
fn from_f64(val: f64) -> Self {
|
|
f16::from_f64(val)
|
|
}
|
|
#[inline]
|
|
fn to_f64(self) -> f64 {
|
|
self.to_f64()
|
|
}
|
|
#[inline]
|
|
fn from_f32(val: f32) -> Self {
|
|
f16::from_f32(val)
|
|
}
|
|
#[inline]
|
|
fn to_f32(self) -> f32 {
|
|
self.to_f32()
|
|
}
|
|
#[inline]
|
|
fn zero() -> Self {
|
|
f16::ZERO
|
|
}
|
|
#[inline]
|
|
fn one() -> Self {
|
|
f16::ONE
|
|
}
|
|
#[inline]
|
|
fn min_value() -> Self {
|
|
f16::MIN
|
|
}
|
|
#[inline]
|
|
fn max_value() -> Self {
|
|
f16::MAX
|
|
}
|
|
#[inline]
|
|
fn nan() -> Self {
|
|
f16::NAN
|
|
}
|
|
#[inline]
|
|
fn infinity() -> Self {
|
|
f16::INFINITY
|
|
}
|
|
#[inline]
|
|
fn neg_infinity() -> Self {
|
|
f16::NEG_INFINITY
|
|
}
|
|
#[inline]
|
|
fn is_nan(self) -> bool {
|
|
self.is_nan()
|
|
}
|
|
#[inline]
|
|
fn is_infinite(self) -> bool {
|
|
self.is_infinite()
|
|
}
|
|
#[inline]
|
|
fn epsilon() -> Self {
|
|
f16::EPSILON
|
|
}
|
|
}
|
|
|
|
impl FloatElement for bf16 {
|
|
#[inline]
|
|
fn name() -> &'static str {
|
|
"bf16"
|
|
}
|
|
#[inline]
|
|
fn bytes() -> usize {
|
|
2
|
|
}
|
|
#[inline]
|
|
fn from_f64(val: f64) -> Self {
|
|
bf16::from_f64(val)
|
|
}
|
|
#[inline]
|
|
fn to_f64(self) -> f64 {
|
|
self.to_f64()
|
|
}
|
|
#[inline]
|
|
fn from_f32(val: f32) -> Self {
|
|
bf16::from_f32(val)
|
|
}
|
|
#[inline]
|
|
fn to_f32(self) -> f32 {
|
|
self.to_f32()
|
|
}
|
|
#[inline]
|
|
fn zero() -> Self {
|
|
bf16::ZERO
|
|
}
|
|
#[inline]
|
|
fn one() -> Self {
|
|
bf16::ONE
|
|
}
|
|
#[inline]
|
|
fn min_value() -> Self {
|
|
bf16::MIN
|
|
}
|
|
#[inline]
|
|
fn max_value() -> Self {
|
|
bf16::MAX
|
|
}
|
|
#[inline]
|
|
fn nan() -> Self {
|
|
bf16::NAN
|
|
}
|
|
#[inline]
|
|
fn infinity() -> Self {
|
|
bf16::INFINITY
|
|
}
|
|
#[inline]
|
|
fn neg_infinity() -> Self {
|
|
bf16::NEG_INFINITY
|
|
}
|
|
#[inline]
|
|
fn is_nan(self) -> bool {
|
|
self.is_nan()
|
|
}
|
|
#[inline]
|
|
fn is_infinite(self) -> bool {
|
|
self.is_infinite()
|
|
}
|
|
#[inline]
|
|
fn epsilon() -> Self {
|
|
bf16::EPSILON
|
|
}
|
|
}
|
|
|
|
/// Trait for integer element types.
|
|
///
|
|
/// Supported types:
|
|
/// - `i8`, `i16`, `i32`, `i64` - Signed integers
|
|
/// - `u8`, `u16`, `u32`, `u64` - Unsigned integers
|
|
pub trait IntElement:
|
|
Clone
|
|
+ Copy
|
|
+ Debug
|
|
+ Default
|
|
+ Display
|
|
+ PartialEq
|
|
+ Eq
|
|
+ PartialOrd
|
|
+ Ord
|
|
+ Send
|
|
+ Sync
|
|
+ Serialize
|
|
+ for<'de> Deserialize<'de>
|
|
+ 'static
|
|
{
|
|
/// The name of this element type.
|
|
fn name() -> &'static str;
|
|
|
|
/// Number of bytes per element.
|
|
fn bytes() -> usize;
|
|
|
|
/// Create from i64.
|
|
fn from_i64(val: i64) -> Self;
|
|
|
|
/// Convert to i64.
|
|
fn to_i64(self) -> i64;
|
|
|
|
/// Zero value.
|
|
fn zero() -> Self;
|
|
|
|
/// One value.
|
|
fn one() -> Self;
|
|
|
|
/// Minimum value.
|
|
fn min_value() -> Self;
|
|
|
|
/// Maximum value.
|
|
fn max_value() -> Self;
|
|
}
|
|
|
|
macro_rules! impl_int_element {
|
|
($t:ty, $name:expr) => {
|
|
impl IntElement for $t {
|
|
fn name() -> &'static str {
|
|
$name
|
|
}
|
|
fn bytes() -> usize {
|
|
std::mem::size_of::<$t>()
|
|
}
|
|
fn from_i64(val: i64) -> Self {
|
|
val as $t
|
|
}
|
|
fn to_i64(self) -> i64 {
|
|
self as i64
|
|
}
|
|
fn zero() -> Self {
|
|
0
|
|
}
|
|
fn one() -> Self {
|
|
1
|
|
}
|
|
fn min_value() -> Self {
|
|
<$t>::MIN
|
|
}
|
|
fn max_value() -> Self {
|
|
<$t>::MAX
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
impl_int_element!(i8, "i8");
|
|
impl_int_element!(i16, "i16");
|
|
impl_int_element!(i32, "i32");
|
|
impl_int_element!(i64, "i64");
|
|
impl_int_element!(u8, "u8");
|
|
impl_int_element!(u16, "u16");
|
|
impl_int_element!(u32, "u32");
|
|
impl_int_element!(u64, "u64");
|
|
|
|
/// Trait for boolean element types.
|
|
///
|
|
/// Used for mask operations and comparisons.
|
|
pub trait BoolElement:
|
|
Clone + Copy + Debug + Default + PartialEq + Eq + Send + Sync + 'static
|
|
{
|
|
/// The name of this element type.
|
|
fn name() -> &'static str;
|
|
|
|
/// True value.
|
|
fn true_val() -> Self;
|
|
|
|
/// False value.
|
|
fn false_val() -> Self;
|
|
|
|
/// Convert to bool.
|
|
fn to_bool(self) -> bool;
|
|
|
|
/// Create from bool.
|
|
fn from_bool(val: bool) -> Self;
|
|
}
|
|
|
|
impl BoolElement for bool {
|
|
#[inline]
|
|
fn name() -> &'static str {
|
|
"bool"
|
|
}
|
|
#[inline]
|
|
fn true_val() -> Self {
|
|
true
|
|
}
|
|
#[inline]
|
|
fn false_val() -> Self {
|
|
false
|
|
}
|
|
#[inline]
|
|
fn to_bool(self) -> bool {
|
|
self
|
|
}
|
|
#[inline]
|
|
fn from_bool(val: bool) -> Self {
|
|
val
|
|
}
|
|
}
|
|
|
|
/// Boolean stored as u8 (useful for GPU where bool isn't native).
|
|
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
|
pub struct BoolU8(pub u8);
|
|
|
|
impl BoolElement for BoolU8 {
|
|
#[inline]
|
|
fn name() -> &'static str {
|
|
"bool_u8"
|
|
}
|
|
#[inline]
|
|
fn true_val() -> Self {
|
|
BoolU8(1)
|
|
}
|
|
#[inline]
|
|
fn false_val() -> Self {
|
|
BoolU8(0)
|
|
}
|
|
#[inline]
|
|
fn to_bool(self) -> bool {
|
|
self.0 != 0
|
|
}
|
|
#[inline]
|
|
fn from_bool(val: bool) -> Self {
|
|
BoolU8(u8::from(val))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_float_element_f32() {
|
|
assert_eq!(f32::name(), "f32");
|
|
assert_eq!(f32::bytes(), 4);
|
|
assert_eq!(f32::from_f64(1.5), 1.5f32);
|
|
assert!(f32::nan().is_nan());
|
|
}
|
|
|
|
#[test]
|
|
fn test_float_element_bf16() {
|
|
assert_eq!(bf16::name(), "bf16");
|
|
assert_eq!(bf16::bytes(), 2);
|
|
let val = bf16::from_f32(1.5);
|
|
assert!((val.to_f32() - 1.5).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_int_element() {
|
|
assert_eq!(i32::name(), "i32");
|
|
assert_eq!(i32::bytes(), 4);
|
|
assert_eq!(i32::from_i64(42), 42i32);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bool_element() {
|
|
assert!(bool::true_val());
|
|
assert!(!bool::false_val());
|
|
assert_eq!(BoolU8::true_val(), BoolU8(1));
|
|
assert!(BoolU8::true_val().to_bool());
|
|
}
|
|
}
|