rust_lstm/
random.rs

1//! Random number generation for weight initialisation, dropout, zoneout and sampling.
2
3use ndarray::{Array2, Ix2, ShapeBuilder};
4use ndarray_rand::RandomExt;
5use rand::rngs::StdRng;
6use rand::{Rng, SeedableRng};
7use rand_distr::Distribution;
8use std::cell::RefCell;
9
10thread_local! {
11    static RNG: RefCell<StdRng> = RefCell::new(StdRng::from_entropy());
12}
13
14/// Reseeds the generator of the current thread.
15///
16/// After `seed(s)`, weight initialisation, dropout and zoneout masks and text sampling
17/// on this thread produce the same values on every run. Each thread has its own
18/// generator, seeded from the operating system until `seed` is called.
19pub fn seed(seed: u64) {
20    RNG.with(|rng| *rng.borrow_mut() = StdRng::seed_from_u64(seed));
21}
22
23pub(crate) fn with_rng<T>(f: impl FnOnce(&mut StdRng) -> T) -> T {
24    RNG.with(|rng| f(&mut rng.borrow_mut()))
25}
26
27pub(crate) fn random_array<Sh, D>(shape: Sh, dist: D) -> Array2<f64>
28where
29    Sh: ShapeBuilder<Dim = Ix2>,
30    D: Distribution<f64>,
31{
32    with_rng(|rng| Array2::random_using(shape, dist, rng))
33}
34
35/// A sample from U[0, 1).
36pub(crate) fn unit() -> f64 {
37    with_rng(|rng| rng.gen())
38}
39
40#[cfg(test)]
41mod tests {
42    use super::*;
43    use ndarray_rand::rand_distr::Uniform;
44
45    #[test]
46    fn same_seed_gives_same_values() {
47        seed(7);
48        let a = random_array((3, 4), Uniform::new(0.0, 1.0));
49        seed(7);
50        let b = random_array((3, 4), Uniform::new(0.0, 1.0));
51        seed(8);
52        let c = random_array((3, 4), Uniform::new(0.0, 1.0));
53        assert_eq!(a, b);
54        assert_ne!(a, c);
55    }
56}
rust_lstm/random.rs

rust_lstm/
random.rs
