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

nuvec_core/
random.rs
