cubecl_common/
rand.rs

1pub use rand::{Rng, SeedableRng, rngs::StdRng};
2
3use rand::distr::StandardUniform;
4use rand::prelude::Distribution;
5
6/// Returns a seeded random number generator using entropy.
7#[cfg(feature = "std")]
8#[inline(always)]
9pub fn get_seeded_rng() -> StdRng {
10    StdRng::from_os_rng()
11}
12
13/// Returns a seeded random number generator using a pre-generated seed.
14#[cfg(not(feature = "std"))]
15#[inline(always)]
16pub fn get_seeded_rng() -> StdRng {
17    const CONST_SEED: u64 = 42;
18    StdRng::seed_from_u64(CONST_SEED)
19}
20
21/// Generates random data from a thread-local RNG.
22#[cfg(feature = "std")]
23#[inline]
24pub fn gen_random<T>() -> T
25where
26    StandardUniform: Distribution<T>,
27{
28    rand::rng().random()
29}
30
31/// Generates random data from a mutex-protected RNG.
32#[cfg(not(feature = "std"))]
33#[inline]
34pub fn gen_random<T>() -> T
35where
36    StandardUniform: Distribution<T>,
37{
38    use crate::stub::Mutex;
39    static RNG: Mutex<Option<StdRng>> = Mutex::new(None);
40    let mut rng = RNG.lock().unwrap();
41    if rng.is_none() {
42        *rng = Some(get_seeded_rng());
43    }
44    rng.as_mut().unwrap().random()
45}