Skip to main content

luma_tensor/
scalar.rs

1//! Scalar values mirroring [`crate::DType`] variants.
2//!
3//! Used for optimizer hyperparameters and other non-tensor data in checkpoint files.
4
5use crate::DType;
6
7/// A typed scalar value corresponding to a [`DType`] variant.
8///
9/// Unlike a 0-d tensor this carries no device, layout, or autograd metadata —
10/// it is just a plain value with a known element type.
11#[derive(Clone, Debug, PartialEq)]
12pub enum Scalar {
13    F32(f32),
14    F64(f64),
15    I32(i32),
16    U32(u32),
17    U8(u8),
18    Bool(bool),
19}
20
21impl Scalar {
22    /// The [`DType`] of this scalar.
23    pub fn dtype(&self) -> DType {
24        match self {
25            Scalar::F32(_) => DType::F32,
26            Scalar::F64(_) => DType::F64,
27            Scalar::I32(_) => DType::I32,
28            Scalar::U32(_) => DType::U32,
29            Scalar::U8(_) => DType::U8,
30            Scalar::Bool(_) => DType::Bool,
31        }
32    }
33
34    /// View as `f64` if the scalar is a float variant.
35    pub fn to_f64(&self) -> Option<f64> {
36        match self {
37            Scalar::F32(v) => Some(*v as f64),
38            Scalar::F64(v) => Some(*v),
39            _ => None,
40        }
41    }
42
43    /// View as `i64` if the scalar is an integer variant.
44    pub fn to_i64(&self) -> Option<i64> {
45        match self {
46            Scalar::I32(v) => Some(*v as i64),
47            Scalar::U32(v) => Some(*v as i64),
48            Scalar::U8(v) => Some(*v as i64),
49            _ => None,
50        }
51    }
52
53    /// View as `bool` if the scalar is the bool variant.
54    pub fn to_bool(&self) -> Option<bool> {
55        match self {
56            Scalar::Bool(v) => Some(*v),
57            _ => None,
58        }
59    }
60}
61
62// ---- Convenience From impls ----
63
64impl From<f64> for Scalar {
65    fn from(v: f64) -> Self {
66        Scalar::F64(v)
67    }
68}
69
70impl From<f32> for Scalar {
71    fn from(v: f32) -> Self {
72        Scalar::F32(v)
73    }
74}
75
76impl From<i64> for Scalar {
77    fn from(v: i64) -> Self {
78        Scalar::I32(v as i32)
79    }
80}
81
82impl From<i32> for Scalar {
83    fn from(v: i32) -> Self {
84        Scalar::I32(v)
85    }
86}
87
88impl From<bool> for Scalar {
89    fn from(v: bool) -> Self {
90        Scalar::Bool(v)
91    }
92}
93
94impl std::fmt::Display for Scalar {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        match self {
97            Scalar::F32(v) => write!(f, "{}", v),
98            Scalar::F64(v) => write!(f, "{}", v),
99            Scalar::I32(v) => write!(f, "{}", v),
100            Scalar::U32(v) => write!(f, "{}", v),
101            Scalar::U8(v) => write!(f, "{}", v),
102            Scalar::Bool(v) => write!(f, "{}", v),
103        }
104    }
105}