nuvec_core/random/
imp.rs

1//! The implementation of the random number functions (see [`crate::random`]).
2
3use std::cell::RefCell;
4use std::sync::OnceLock;
5use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
6
7use crate::ops::imp::{add, divide, erfinv, minimum, multiply, slice, subtract};
8use crate::primitive::Primitive;
9use crate::{Array, DType, Error, Result, SliceArg, bf16, f16};
10
11/// The body of [`crate::random::key`].
12pub fn key(seed: u64) -> Array {
13    let data = [(seed >> 32) as u32, seed as u32];
14    crate::creation::from_data(&data, &[2]).unwrap_or_else(|e| Array::failed(e, DType::U32, &[2]))
15}
16
17fn time_seed() -> u64 {
18    static SEED: OnceLock<u64> = OnceLock::new();
19    *SEED.get_or_init(|| {
20        // Miri isolates the clock: a fixed seed there (reproducible runs too).
21        if cfg!(miri) {
22            return 0x5eed;
23        }
24        std::time::SystemTime::now()
25            .duration_since(std::time::UNIX_EPOCH)
26            .map_or(0, |d| d.as_millis() as u64)
27    })
28}
29
30/// The last [`seed`], once one was given (else the time seed applies).
31static PROCESS_SEED: AtomicU64 = AtomicU64::new(0);
32static SEEDED: AtomicBool = AtomicBool::new(false);
33/// Bumped by [`seed`]: sequences derived before are stale.
34static EPOCH: AtomicU64 = AtomicU64::new(0);
35/// The next thread ordinal (reset by [`seed`]).
36static THREAD_COUNTER: AtomicU64 = AtomicU64::new(0);
37
38fn process_seed() -> u64 {
39    match SEEDED.load(Ordering::Acquire) {
40        true => PROCESS_SEED.load(Ordering::Acquire),
41        false => time_seed(),
42    }
43}
44
45/// One step of SplitMix64: spreads nearby seeds and ordinals over unrelated keys.
46fn splitmix64(x: u64) -> u64 {
47    let mut z = x.wrapping_add(0x9E37_79B9_7F4A_7C15);
48    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
49    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
50    z ^ (z >> 31)
51}
52
53/// The first key of the sequence of thread `ordinal` (see the module documentation).
54fn derived_key(seed: u64, ordinal: u64) -> Array {
55    match ordinal {
56        0 => key(seed),
57        n => key(splitmix64(seed ^ n)),
58    }
59}
60
61/// A thread's key sequence.
62struct Sequence {
63    key: Array,
64    /// The [`EPOCH`] it was derived in.
65    epoch: u64,
66    /// Set by this thread's own [`seed`] call: kept when another thread seeds.
67    seeded_here: bool,
68}
69
70thread_local!(static SEQUENCE: RefCell<Option<Sequence>> = const { RefCell::new(None) });
71
72/// The body of [`crate::random::seed`].
73pub fn seed(seed: u64) {
74    PROCESS_SEED.store(seed, Ordering::Release);
75    SEEDED.store(true, Ordering::Release);
76    THREAD_COUNTER.store(1, Ordering::Release);
77    let epoch = EPOCH.fetch_add(1, Ordering::AcqRel) + 1;
78    SEQUENCE.set(Some(Sequence {
79        key: derived_key(seed, 0),
80        epoch,
81        seeded_here: true,
82    }));
83}
84
85/// The body of [`crate::random::next_key`].
86pub(crate) fn next_key() -> Result<Array> {
87    let current = sequence_key();
88    let (keep, out) = split_two(&current)?;
89    SEQUENCE.with_borrow_mut(|s| {
90        if let Some(s) = s {
91            s.key = keep;
92        }
93    });
94    Ok(out)
95}
96
97fn key_or_next(key: Option<&Array>) -> Result<Array> {
98    match key {
99        Some(k) => Ok(k.clone()),
100        None => next_key(),
101    }
102}
103
104/// The body of [`crate::random::bits`].
105pub(crate) fn bits(shape: &[usize], width: usize, key: Option<&Array>) -> Result<Array> {
106    let key = key_or_next(key)?;
107    if key.dtype() != DType::U32 || key.dims() != [2] {
108        return Err(Error::InvalidArgument(format!(
109            "bits: expected a uint32 key of shape [2], got {} {:?}",
110            key.dtype(),
111            key.dims()
112        )));
113    }
114    let width = u8::try_from(width)
115        .map_err(|_| Error::InvalidArgument(format!("bits: width {width} not in {{1, 2, 4}}")))?;
116    let primitive = Primitive::RandomBits {
117        shape: shape.to_vec(),
118        width,
119    };
120    Ok(Array::from_op(primitive, vec![key])?.swap_remove(0))
121}
122
123/// The body of [`crate::random::split`].
124pub(crate) fn split(key: &Array, num: usize) -> Result<Array> {
125    bits(&[num, 2], 4, Some(key))
126}
127
128/// The body of [`crate::random::split_two`].
129pub(crate) fn split_two(key: &Array) -> Result<(Array, Array)> {
130    let both = split(key, 2)?;
131    let row = |i: isize| {
132        let args = [SliceArg::new(Some(i), Some(i + 1), 1), SliceArg::full()];
133        slice(&both, &args)?.try_reshape(&[2])
134    };
135    Ok((row(0)?, row(1)?))
136}
137
138/// The largest value below 1 in `dtype`, as float32 (samples stay in `[low, high)`).
139fn below_one(dtype: DType) -> Result<f32> {
140    Ok(match dtype {
141        DType::F32 => f32::from_bits(1.0f32.to_bits() - 1),
142        DType::F16 => f16::from_bits(f16::ONE.to_bits() - 1).to_f32(),
143        DType::BF16 => bf16::from_bits(bf16::ONE.to_bits() - 1).to_f32(),
144        _ => {
145            return Err(Error::unsupported_dtype("uniform", dtype));
146        }
147    })
148}
149
150/// The body of [`crate::random::uniform`].
151pub(crate) fn uniform(
152    low: f64,
153    high: f64,
154    shape: &[usize],
155    dtype: DType,
156    key: Option<&Array>,
157) -> Result<Array> {
158    let upper = below_one(dtype)?;
159    let lo = Array::from(low).try_astype(dtype)?;
160    let hi = Array::from(high).try_astype(dtype)?;
161    let range = subtract(&hi, &lo)?;
162    let out = divide(&bits(shape, 4, key)?, Array::from(u32::MAX as f32))?;
163    let out = minimum(&out, Array::from(upper))?.try_astype(dtype)?;
164    add(multiply(&range, &out)?, &lo)
165}
166
167/// The body of [`crate::random::normal`].
168pub(crate) fn normal(
169    shape: &[usize],
170    dtype: DType,
171    loc: Option<f64>,
172    scale: Option<f64>,
173    key: Option<&Array>,
174) -> Result<Array> {
175    if !matches!(dtype, DType::F16 | DType::BF16 | DType::F32 | DType::F64) {
176        return Err(Error::unsupported_dtype("normal", dtype));
177    }
178    let low = f64::from(f32::from_bits((-1.0f32).to_bits() - 1));
179    let samples = uniform(low, 1.0, shape, DType::F32, key)?;
180    let mut applied = Array::from(2f64.sqrt()).try_astype(dtype)?;
181    if let Some(scale) = scale {
182        applied = multiply(&applied, Array::from(scale).try_astype(dtype)?)?;
183    }
184    let samples = multiply(&applied, erfinv(&samples)?.try_astype(dtype)?)?;
185    match loc {
186        Some(loc) => add(Array::from(loc).try_astype(dtype)?, &samples),
187        None => Ok(samples),
188    }
189}
190
191/// The body of [`crate::random::gumbel`].
192pub(crate) fn gumbel(shape: &[usize], dtype: DType, key: Option<&Array>) -> Result<Array> {
193    let u = uniform(0.0, 1.0, shape, dtype, key)?;
194    crate::ops::imp::negative(&crate::ops::imp::log(&crate::ops::imp::negative(
195        &crate::ops::imp::log(&u)?,
196    )?)?)
197}
198
199/// The body of [`crate::random::categorical`].
200pub(crate) fn categorical(logits: &Array, axis: isize, key: Option<&Array>) -> Result<Array> {
201    let noise = gumbel(logits.dims(), logits.dtype(), key)?;
202    crate::ops::imp::argmax(&add(&noise, logits)?, Some(axis), false)
203}
204
205/// The body of [`crate::random::randint`].
206pub(crate) fn randint(
207    low: i64,
208    high: i64,
209    shape: &[usize],
210    dtype: DType,
211    key: Option<&Array>,
212) -> Result<Array> {
213    if dtype.category() == crate::Category::Float || dtype == DType::Complex64 {
214        return Err(Error::unsupported_dtype("randint", dtype));
215    }
216    let u = uniform(low as f64, high as f64, shape, DType::F32, key)?;
217    let u = crate::ops::imp::maximum(&u, Array::from(low as f32))?;
218    // floor(u): truncation, one less where it rounded a negative value up.
219    let t = u.try_astype(DType::I64)?;
220    let up = crate::ops::imp::less(&u, &t.try_astype(DType::F32)?)?;
221    subtract(&t, up.try_astype(DType::I64)?)?.try_astype(dtype)
222}
223
224/// The body of [`crate::random::bernoulli`].
225pub(crate) fn bernoulli(p: f64, shape: &[usize], key: Option<&Array>) -> Result<Array> {
226    crate::ops::imp::less(&uniform(0.0, 1.0, shape, DType::F32, key)?, p)
227}
228
229/// The current key of this thread's sequence (derived first if needed), without advancing
230/// it: the body of [`crate::random::sequence_key`].
231pub(crate) fn sequence_key() -> Array {
232    SEQUENCE.with_borrow_mut(|s| {
233        let epoch = EPOCH.load(Ordering::Acquire);
234        let stale = s
235            .as_ref()
236            .is_none_or(|s| !s.seeded_here && s.epoch != epoch);
237        if stale {
238            let ordinal = THREAD_COUNTER.fetch_add(1, Ordering::AcqRel);
239            *s = Some(Sequence {
240                key: derived_key(process_seed(), ordinal),
241                epoch,
242                seeded_here: false,
243            });
244        }
245        s.as_ref().expect("set above").key.clone()
246    })
247}
248
249/// Replaces the current key of this thread's sequence: the body of
250/// [`crate::random::set_sequence_key`].
251pub(crate) fn set_sequence_key(key: Array) {
252    let _ = sequence_key(); // derive the sequence first, so its epoch is current
253    SEQUENCE.with_borrow_mut(|s| {
254        if let Some(s) = s {
255            s.key = key;
256        }
257    });
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263    use crate::kernels::threefry2x32;
264
265    #[test]
266    fn threefry_known_answer() {
267        // Random123's known-answer vector for threefry2x32-20 with a zero key and counter.
268        assert_eq!(threefry2x32((0, 0), (0, 0)), (0x6b20_0159, 0x99ba_4efe));
269    }
270
271    #[test]
272    fn keys_and_layout() {
273        assert_eq!(key(1 << 32 | 7).to_vec::<u32>().unwrap(), [1, 7]);
274        // An odd word count and a partial final word use the tail paths.
275        let k = key(3);
276        for (n, width) in [(5, 4), (3, 2), (7, 1), (0, 4)] {
277            let b = bits(&[n], width, Some(&k)).unwrap();
278            b.eval().unwrap();
279            assert_eq!(b.size(), n);
280        }
281        assert!(bits(&[2], 3, Some(&k)).is_err());
282        assert!(bits(&[2], 4, Some(&Array::from(1u32))).is_err());
283    }
284
285    #[test]
286    fn sequences_are_per_thread_and_reproducible() {
287        seed(42);
288        let a = uniform(0.0, 1.0, &[4], DType::F32, None)
289            .unwrap()
290            .to_vec::<f32>()
291            .unwrap();
292        seed(42);
293        let b = uniform(0.0, 1.0, &[4], DType::F32, None)
294            .unwrap()
295            .to_vec::<f32>()
296            .unwrap();
297        assert_eq!(a, b);
298        assert!(a.iter().all(|&x| (0.0..1.0).contains(&x)));
299    }
300}
nuvec_core/random/imp.rs

nuvec_core/random/
imp.rs
