nuvec_core/
scalar.rs

1//! Weak-typed scalar operands, as in MLX's Python layer (`python/src/utils.cpp`).
2//!
3//! A Rust number combined with an array adapts to the array: `&f16_array * 2.0` stays
4//! float16, `&int8_array + 1` stays int8 (and `+ 300` is an overflow error), `&bool_array + 1`
5//! is int32, `&int_array * 0.5` is float32. Every Rust scalar is weak, including typed
6//! literals such as `2.0f32`; use [`Array::from`] for a strongly typed scalar.
7
8use crate::cast::Canon;
9use crate::{Array, Category, Complex64, DType, Error, Result, bf16, f16};
10
11/// A number whose dtype is decided by the array it is combined with.
12#[derive(Debug, Clone, Copy, PartialEq)]
13pub enum Scalar {
14    /// A boolean.
15    Bool(bool),
16    /// An integer (exact for every Rust integer type).
17    Int(i128),
18    /// A real floating-point number.
19    Float(f64),
20    /// A complex number.
21    Complex(Complex64),
22}
23
24/// An operand of a binary operation: an array, or a weak scalar.
25#[derive(Debug, Clone)]
26pub enum ScalarOrArray {
27    /// An array; its dtype is kept.
28    Array(Array),
29    /// A weak scalar; its dtype follows the other operand.
30    Scalar(Scalar),
31}
32
33impl From<Array> for ScalarOrArray {
34    fn from(a: Array) -> Self {
35        Self::Array(a)
36    }
37}
38
39impl From<&Array> for ScalarOrArray {
40    fn from(a: &Array) -> Self {
41        Self::Array(a.clone())
42    }
43}
44
45impl From<Scalar> for ScalarOrArray {
46    fn from(s: Scalar) -> Self {
47        Self::Scalar(s)
48    }
49}
50
51/// A Rust number usable as a weak operand of `+ - * /` on the right of an array.
52///
53/// Implemented for `bool`, every integer type, `f16`, `bf16`, `f32`, `f64` and `Complex64`;
54/// a single generic operator impl over this trait lets literals such as `2.0` and `1` keep
55/// their usual defaults.
56pub trait ScalarOperand: Into<Scalar> {}
57
58macro_rules! from_scalar {
59    ($($t:ty => |$x:ident| $e:expr),* $(,)?) => {$(
60        impl ScalarOperand for $t {}
61        impl From<$t> for Scalar {
62            fn from($x: $t) -> Self {
63                $e
64            }
65        }
66        impl From<$t> for ScalarOrArray {
67            fn from(x: $t) -> Self {
68                Self::Scalar(Scalar::from(x))
69            }
70        }
71    )*};
72}
73
74from_scalar! {
75    bool => |x| Scalar::Bool(x),
76    i8 => |x| Scalar::Int(x.into()),
77    i16 => |x| Scalar::Int(x.into()),
78    i32 => |x| Scalar::Int(x.into()),
79    i64 => |x| Scalar::Int(x.into()),
80    u8 => |x| Scalar::Int(x.into()),
81    u16 => |x| Scalar::Int(x.into()),
82    u32 => |x| Scalar::Int(x.into()),
83    u64 => |x| Scalar::Int(x.into()),
84    f16 => |x| Scalar::Float(x.to_f64()),
85    bf16 => |x| Scalar::Float(x.to_f64()),
86    f32 => |x| Scalar::Float(x.into()),
87    f64 => |x| Scalar::Float(x),
88    Complex64 => |x| Scalar::Complex(x),
89}
90
91/// The range of an integer dtype.
92fn int_range(d: DType) -> (i128, i128) {
93    match d {
94        DType::U8 => (0, u8::MAX.into()),
95        DType::U16 => (0, u16::MAX.into()),
96        DType::U32 => (0, u32::MAX.into()),
97        DType::U64 => (0, u64::MAX.into()),
98        DType::I8 => (i8::MIN.into(), i8::MAX.into()),
99        DType::I16 => (i16::MIN.into(), i16::MAX.into()),
100        DType::I32 => (i32::MIN.into(), i32::MAX.into()),
101        _ => (i64::MIN.into(), i64::MAX.into()),
102    }
103}
104
105impl Scalar {
106    /// The dtype of a scalar on its own: bool, int32 (int64, then uint64, when out of
107    /// range), float32 or complex64.
108    pub fn default_dtype(self) -> DType {
109        match self {
110            Self::Bool(_) => DType::Bool,
111            Self::Int(v) if i32::try_from(v).is_ok() => DType::I32,
112            Self::Int(v) if i64::try_from(v).is_ok() => DType::I64,
113            Self::Int(_) => DType::U64,
114            Self::Float(_) => DType::F32,
115            Self::Complex(_) => DType::Complex64,
116        }
117    }
118
119    /// The dtype this scalar takes next to an array of dtype `like` (MLX's `to_array`).
120    fn dtype_like(self, like: DType) -> Result<DType> {
121        Ok(match self {
122            Self::Bool(_) => like,
123            Self::Int(v) => match like.category() {
124                // bool is always promoted.
125                Category::Bool => Self::Int(v).default_dtype(),
126                Category::Signed | Category::Unsigned => {
127                    let (lo, hi) = int_range(like);
128                    if !(lo..=hi).contains(&v) {
129                        return Err(Error::InvalidArgument(format!(
130                            "converting {v} to {like} would overflow"
131                        )));
132                    }
133                    like
134                }
135                _ => like,
136            },
137            Self::Float(_) if like.category() == Category::Float => like,
138            Self::Float(_) => DType::F32,
139            Self::Complex(_) => DType::Complex64,
140        })
141    }
142
143    pub(crate) fn canon(self) -> Canon {
144        match self {
145            Self::Bool(b) => Canon::Int(i128::from(b)),
146            Self::Int(v) => Canon::Int(v),
147            Self::Float(f) => Canon::Float(f),
148            Self::Complex(z) => Canon::Complex(z),
149        }
150    }
151
152    /// A 0-D array holding this scalar, typed relative to `like` (or by default).
153    pub(crate) fn to_array(self, like: Option<DType>) -> Result<Array> {
154        let dtype = match like {
155            Some(like) => self.dtype_like(like)?,
156            None => self.default_dtype(),
157        };
158        crate::ops::imp::filled(&[], dtype, |_| self.canon())
159    }
160}
161
162impl ScalarOrArray {
163    /// This operand as an array: weak scalars take a dtype relative to `like`.
164    pub(crate) fn into_array(self, like: Option<DType>) -> Result<Array> {
165        match self {
166            Self::Array(a) => Ok(a),
167            Self::Scalar(s) => s.to_array(like),
168        }
169    }
170
171    fn dtype(&self) -> Option<DType> {
172        match self {
173            Self::Array(a) => Some(a.dtype()),
174            Self::Scalar(_) => None,
175        }
176    }
177}
178
179/// Both operands as arrays, following MLX's `to_arrays`: arrays keep their dtype, a weak
180/// scalar adapts to the other operand, and two scalars take their default dtypes.
181/// An input that has already failed reports its error before any scalar conversion.
182pub(crate) fn resolve(a: ScalarOrArray, b: ScalarOrArray) -> Result<(Array, Array)> {
183    for x in [&a, &b] {
184        if let ScalarOrArray::Array(x) = x
185            && let Some(err) = x.slot_error()
186        {
187            return Err(err);
188        }
189    }
190    let (da, db) = (a.dtype(), b.dtype());
191    Ok((a.into_array(db)?, b.into_array(da)?))
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    fn like(s: impl Into<Scalar>, d: DType) -> Result<DType> {
199        s.into().dtype_like(d)
200    }
201
202    #[test]
203    fn mlx_weak_typing_rules() {
204        assert_eq!(like(2.0, DType::F16).unwrap(), DType::F16);
205        assert_eq!(like(2.0, DType::I32).unwrap(), DType::F32);
206        assert_eq!(like(2.0, DType::F64).unwrap(), DType::F64);
207        assert_eq!(like(2, DType::I8).unwrap(), DType::I8);
208        assert_eq!(like(2, DType::F16).unwrap(), DType::F16);
209        assert_eq!(like(2, DType::Bool).unwrap(), DType::I32);
210        assert_eq!(like(true, DType::U16).unwrap(), DType::U16);
211        assert_eq!(
212            like(Complex64::new(1.0, 1.0), DType::F32).unwrap(),
213            DType::Complex64
214        );
215        assert!(matches!(
216            like(300, DType::I8),
217            Err(Error::InvalidArgument(_))
218        ));
219        assert!(matches!(
220            like(-1, DType::U32),
221            Err(Error::InvalidArgument(_))
222        ));
223        assert_eq!(like(u64::MAX, DType::U64).unwrap(), DType::U64);
224    }
225
226    #[test]
227    fn defaults() {
228        assert_eq!(Scalar::from(1).default_dtype(), DType::I32);
229        assert_eq!(Scalar::from(1i64 << 40).default_dtype(), DType::I64);
230        assert_eq!(Scalar::from(u64::MAX).default_dtype(), DType::U64);
231        assert_eq!(Scalar::from(1.0f32).default_dtype(), DType::F32);
232        assert_eq!(Scalar::from(false).default_dtype(), DType::Bool);
233    }
234}
nuvec_core/scalar.rs

nuvec_core/
scalar.rs
