nuvec_core/
random.rs

1//! Random numbers from explicit keys or a per-thread key sequence, bit for bit as in MLX.
2//!
3//! A key is a uint32 array of shape `[2]`. [`bits`] hashes counters with threefry2x32
4//! using MLX's layout, and [`uniform`] and [`normal`] are composed from operations
5//! exactly as in MLX's `random.cpp`. For the same seed, bits and uniform samples equal MLX's.
6//! Normal samples agree within a few ulps: `erfinv` needs a logarithm, whose last bits
7//! differ between implementations.
8//!
9//! Without a key, functions draw the next key from this thread's sequence, which starts from
10//! a per-process time seed until [`seed`] is called (as in MLX).
11
12use std::cell::RefCell;
13use std::sync::OnceLock;
14
15use crate::ops::{add, divide, erfinv, minimum, multiply, slice, subtract};
16use crate::primitive::Primitive;
17use crate::{Array, DType, Error, Result, SliceArg, bf16, f16};
18
19/// A key for `seed`: `[seed >> 32, seed & 0xffffffff]`.
20pub fn key(seed: u64) -> Array {
21    Array::from_slice(&[(seed >> 32) as u32, seed as u32], &[2]).expect("two elements")
22}
23
24fn time_seed() -> u64 {
25    static SEED: OnceLock<u64> = OnceLock::new();
26    *SEED.get_or_init(|| {
27        std::time::SystemTime::now()
28            .duration_since(std::time::UNIX_EPOCH)
29            .map_or(0, |d| d.as_millis() as u64)
30    })
31}
32
33thread_local!(static SEQUENCE: RefCell<Option<Array>> = const { RefCell::new(None) });
34
35/// Restarts this thread's key sequence from `seed`.
36pub fn seed(seed: u64) {
37    SEQUENCE.set(Some(key(seed)));
38}
39
40/// The next key of this thread's sequence (`split` the current key; keep the first half),
41/// the key that functions called without one use. Useful to pass the sequence's state into
42/// a compiled function, where keys drawn inside would be frozen at the first trace.
43pub fn next_key() -> Result<Array> {
44    let current = SEQUENCE.with_borrow_mut(|k| k.get_or_insert_with(|| key(time_seed())).clone());
45    let (keep, out) = split_pair(&current)?;
46    SEQUENCE.set(Some(keep));
47    Ok(out)
48}
49
50fn key_or_next(key: Option<&Array>) -> Result<Array> {
51    match key {
52        Some(k) => Ok(k.clone()),
53        None => next_key(),
54    }
55}
56
57/// Random bits of `shape`, `width` bytes per element (1, 2 or 4 → uint8, uint16, uint32).
58pub fn bits(shape: &[usize], width: usize, key: Option<&Array>) -> Result<Array> {
59    let key = key_or_next(key)?;
60    if key.dtype() != DType::U32 || key.shape() != [2] {
61        return Err(Error::InvalidArgument(format!(
62            "bits: expected a uint32 key of shape [2], got {} {:?}",
63            key.dtype(),
64            key.shape()
65        )));
66    }
67    let width = u8::try_from(width)
68        .map_err(|_| Error::InvalidArgument(format!("bits: width {width} not in {{1, 2, 4}}")))?;
69    let primitive = Primitive::RandomBits {
70        shape: shape.to_vec(),
71        width,
72    };
73    Ok(Array::from_op(primitive, vec![key])?.swap_remove(0))
74}
75
76/// `num` new keys from `key`, as an array of shape `[num, 2]`.
77pub fn split(key: &Array, num: usize) -> Result<Array> {
78    bits(&[num, 2], 4, Some(key))
79}
80
81/// Two new keys from `key`.
82pub fn split_pair(key: &Array) -> Result<(Array, Array)> {
83    let both = split(key, 2)?;
84    let row = |i: isize| {
85        let args = [SliceArg::new(Some(i), Some(i + 1), 1), SliceArg::full()];
86        slice(&both, &args)?.reshape(&[2])
87    };
88    Ok((row(0)?, row(1)?))
89}
90
91/// The largest value below 1 in `dtype`, as float32 (samples stay in `[low, high)`).
92fn below_one(dtype: DType) -> Result<f32> {
93    Ok(match dtype {
94        DType::F32 => f32::from_bits(1.0f32.to_bits() - 1),
95        DType::F16 => f16::from_bits(f16::ONE.to_bits() - 1).to_f32(),
96        DType::BF16 => bf16::from_bits(bf16::ONE.to_bits() - 1).to_f32(),
97        _ => {
98            return Err(Error::UnsupportedDType {
99                op: "uniform",
100                dtype,
101            });
102        }
103    })
104}
105
106/// Uniform samples in `[low, high)` of `shape` and floating-point `dtype` (float16,
107/// bfloat16 or float32), as MLX: 32 random bits divided by `u32::MAX`, clamped below 1.
108pub fn uniform(
109    low: f64,
110    high: f64,
111    shape: &[usize],
112    dtype: DType,
113    key: Option<&Array>,
114) -> Result<Array> {
115    let upper = below_one(dtype)?;
116    let lo = Array::scalar(low).astype(dtype)?;
117    let hi = Array::scalar(high).astype(dtype)?;
118    let range = subtract(&hi, &lo)?;
119    let out = divide(&bits(shape, 4, key)?, Array::scalar(u32::MAX as f32))?;
120    let out = minimum(&out, Array::scalar(upper))?.astype(dtype)?;
121    add(multiply(&range, &out)?, &lo)
122}
123
124/// Normal samples (mean `loc`, standard deviation `scale`) of `shape` and floating-point
125/// `dtype`, as MLX: `√2 · erfinv(u)` for `u` uniform in `(-1, 1)`.
126pub fn normal(
127    shape: &[usize],
128    dtype: DType,
129    loc: Option<f64>,
130    scale: Option<f64>,
131    key: Option<&Array>,
132) -> Result<Array> {
133    if !matches!(dtype, DType::F16 | DType::BF16 | DType::F32 | DType::F64) {
134        return Err(Error::UnsupportedDType {
135            op: "normal",
136            dtype,
137        });
138    }
139    let low = f64::from(f32::from_bits((-1.0f32).to_bits() - 1));
140    let samples = uniform(low, 1.0, shape, DType::F32, key)?;
141    let mut applied = Array::scalar(2f64.sqrt()).astype(dtype)?;
142    if let Some(scale) = scale {
143        applied = multiply(&applied, Array::scalar(scale).astype(dtype)?)?;
144    }
145    let samples = multiply(&applied, erfinv(&samples)?.astype(dtype)?)?;
146    match loc {
147        Some(loc) => add(Array::scalar(loc).astype(dtype)?, &samples),
148        None => Ok(samples),
149    }
150}
151
152/// Gumbel samples of `shape` and floating-point `dtype`, as MLX: `−log(−log(u))` for `u`
153/// uniform in `[0, 1)`.
154pub fn gumbel(shape: &[usize], dtype: DType, key: Option<&Array>) -> Result<Array> {
155    let u = uniform(0.0, 1.0, shape, dtype, key)?;
156    crate::ops::negative(&crate::ops::log(&crate::ops::negative(&crate::ops::log(
157        &u,
158    )?)?)?)
159}
160
161/// Indices (uint32) drawn from the categorical distributions whose unnormalised log
162/// probabilities run along `axis` of `logits`, as MLX: the argmax of `logits` plus Gumbel
163/// noise. The result has `logits`' shape without `axis`.
164///
165/// ```
166/// use nuvec_core::{Array, Result, random};
167///
168/// # fn main() -> Result<()> {
169/// // One certain outcome per row.
170/// let logits = Array::from_slice(&[0.0f32, -1e9, -1e9, -1e9, -1e9, 0.0], &[2, 3])?;
171/// let draws = random::categorical(&logits, -1, Some(&random::key(7)))?;
172/// assert_eq!(draws.to_vec::<u32>()?, [0, 2]);
173/// # Ok(())
174/// # }
175/// ```
176pub fn categorical(logits: &Array, axis: isize, key: Option<&Array>) -> Result<Array> {
177    let noise = gumbel(logits.shape(), logits.dtype(), key)?;
178    crate::ops::argmax(&add(&noise, logits)?, Some(axis), false)
179}
180
181/// Integers in `[low, high)` of `shape` and integer `dtype`, as MLX: uniform floats in
182/// `[low, high)` rounded down.
183pub fn randint(
184    low: i64,
185    high: i64,
186    shape: &[usize],
187    dtype: DType,
188    key: Option<&Array>,
189) -> Result<Array> {
190    if dtype.category() == crate::Category::Float || dtype == DType::Complex64 {
191        return Err(Error::UnsupportedDType {
192            op: "randint",
193            dtype,
194        });
195    }
196    let u = uniform(low as f64, high as f64, shape, DType::F32, key)?;
197    let u = crate::ops::maximum(&u, Array::scalar(low as f32))?;
198    // floor(u): truncation, one less where it rounded a negative value up.
199    let t = u.astype(DType::I64)?;
200    let up = crate::ops::less(&u, &t.astype(DType::F32)?)?;
201    subtract(&t, up.astype(DType::I64)?)?.astype(dtype)
202}
203
204/// Booleans of `shape` that are true with probability `p`, as MLX: `uniform < p`.
205pub fn bernoulli(p: f64, shape: &[usize], key: Option<&Array>) -> Result<Array> {
206    crate::ops::less(&uniform(0.0, 1.0, shape, DType::F32, key)?, p)
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212    use crate::kernels::threefry2x32;
213
214    #[test]
215    fn threefry_known_answer() {
216        // Random123's known-answer vector for threefry2x32-20 with a zero key and counter.
217        assert_eq!(threefry2x32((0, 0), (0, 0)), (0x6b20_0159, 0x99ba_4efe));
218    }
219
220    #[test]
221    fn keys_and_layout() {
222        assert_eq!(key(1 << 32 | 7).to_vec::<u32>().unwrap(), [1, 7]);
223        // An odd word count and a partial final word use the tail paths.
224        let k = key(3);
225        for (n, width) in [(5, 4), (3, 2), (7, 1), (0, 4)] {
226            let b = bits(&[n], width, Some(&k)).unwrap();
227            b.eval().unwrap();
228            assert_eq!(b.size(), n);
229        }
230        assert!(bits(&[2], 3, Some(&k)).is_err());
231        assert!(bits(&[2], 4, Some(&Array::scalar(1u32))).is_err());
232    }
233
234    #[test]
235    fn sequences_are_per_thread_and_reproducible() {
236        seed(42);
237        let a = uniform(0.0, 1.0, &[4], DType::F32, None)
238            .unwrap()
239            .to_vec::<f32>()
240            .unwrap();
241        seed(42);
242        let b = uniform(0.0, 1.0, &[4], DType::F32, None)
243            .unwrap()
244            .to_vec::<f32>()
245            .unwrap();
246        assert_eq!(a, b);
247        assert!(a.iter().all(|&x| (0.0..1.0).contains(&x)));
248    }
249}
nuvec_core/random.rs

nuvec_core/
random.rs
