fuel-tensor 0.19.1

Fuel's user-facing Tensor handle and Device abstraction — the central, near-universal consumer surface of the former fuel-core (board #109; CireSnave: "Move Tensor into fuel-tensor. We can worry about a semantic split later.")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Types for elements that can be stored and manipulated using tensors.
//!
//! The [`DType`] enum and its inherent methods, Display, FromStr, DTypeParseError,
//! and safetensors interop are all defined in `fuel-core-types` and re-exported here.
//!
//! [`WithDType`] is a local marker subtrait of `fuel_ir::dtype::WithDType`.
//! Keeping it local lets Rust's coherence checker prove disjointness of generic impls
//! (e.g. `NdArray for S` vs `NdArray for Vec<S>`).  All methods are inherited from
//! the upstream trait.
#![allow(clippy::redundant_closure_call)]

// Re-export DType and related types from fuel-core-types.
pub use fuel_ir::dtype::{DType, DTypeParseError};

/// Local marker subtrait — inherits every method from
/// `fuel_ir::dtype::WithDType` but is defined in this crate so that
/// the orphan/coherence rules treat it as a local trait.
///
/// Also requires `fuel_ir::HostDType` (the host-buffer conversion methods split
/// off `WithDType` in B0.4): every `BuiltinDType` here is host-storable, so this
/// keeps `crate::WithDType`-bounded fuel-core code (NdArray, tensor `to_vec*` /
/// `from_*`) able to call `to_cpu_storage` / `cpu_storage_as_slice` unchanged.
pub trait WithDType: fuel_ir::dtype::WithDType + fuel_ir::HostDType {}

macro_rules! with_dtype {
    ($ty:ty, $dtype:ident, $from_f64:expr, $to_f64:expr) => {
        impl WithDType for $ty {}
    };
}
use float8::F8E4M3 as f8e4m3;
use half::{bf16, f16};

with_dtype!(u8, U8, |v: f64| v as u8, |v: u8| v as f64);
with_dtype!(i8, I8, |v: f64| v as i8, |v: i8| v as f64);
with_dtype!(u32, U32, |v: f64| v as u32, |v: u32| v as f64);
with_dtype!(i16, I16, |v: f64| v as i16, |v: i16| v as f64);
with_dtype!(i32, I32, |v: f64| v as i32, |v: i32| v as f64);
with_dtype!(i64, I64, |v: f64| v as i64, |v: i64| v as f64);
with_dtype!(f16, F16, f16::from_f64, f16::to_f64);
with_dtype!(bf16, BF16, bf16::from_f64, bf16::to_f64);
with_dtype!(f32, F32, |v: f64| v as f32, |v: f32| v as f64);
with_dtype!(f64, F64, |v: f64| v, |v: f64| v);
with_dtype!(f8e4m3, F8E4M3, f8e4m3::from_f64, |v: f8e4m3| v.to_f64());

pub trait IntDType: WithDType + num_traits::Bounded {
    fn is_true(&self) -> bool;
    fn as_usize(&self) -> usize;
}

impl IntDType for i64 {
    fn is_true(&self) -> bool {
        *self != 0
    }
    fn as_usize(&self) -> usize {
        *self as usize
    }
}

impl IntDType for u32 {
    fn is_true(&self) -> bool {
        *self != 0
    }
    fn as_usize(&self) -> usize {
        *self as usize
    }
}

impl IntDType for u8 {
    fn is_true(&self) -> bool {
        *self != 0
    }
    fn as_usize(&self) -> usize {
        *self as usize
    }
}

impl IntDType for i8 {
    fn is_true(&self) -> bool {
        *self != 0
    }
    fn as_usize(&self) -> usize {
        *self as usize
    }
}

impl IntDType for i16 {
    fn is_true(&self) -> bool {
        *self != 0
    }
    fn as_usize(&self) -> usize {
        *self as usize
    }
}

impl IntDType for i32 {
    fn is_true(&self) -> bool {
        *self != 0
    }
    fn as_usize(&self) -> usize {
        *self as usize
    }
}

pub trait FloatDType: WithDType {}

impl FloatDType for f16 {}
impl FloatDType for bf16 {}
impl FloatDType for f32 {}
impl FloatDType for f64 {}
impl FloatDType for f8e4m3 {}