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

nuvec_core/random/
imp.rs
