dataloader_rs/sampler/
random.rs

1use std::hash::{BuildHasher, Hasher, RandomState};
2
3use crate::sampler::Sampler;
4
5/// Yields a uniformly shuffled permutation of indices every epoch.
6///
7/// Uses the inside-out Fisher-Yates algorithm with `fastrand` (wyrand):
8/// a single-multiply PRNG that is ~30% faster than Xoshiro256++ for this
9/// workload. The inside-out variant builds and shuffles in a single forward
10/// pass — no pre-fill, and random accesses stay within the already-written
11/// prefix `[0..i]` for better cache behaviour.
12#[derive(Debug)]
13pub struct RandomSampler {
14    rng: fastrand::Rng,
15}
16
17impl RandomSampler {
18    /// Deterministic seed. Use this for reproducible experiments.
19    pub fn new(seed: u64) -> Self {
20        Self {
21            rng: fastrand::Rng::with_seed(seed),
22        }
23    }
24
25    /// Seed from OS entropy. Use this when reproducibility is not required.
26    pub fn from_entropy() -> Self {
27        // std draws `RandomState`'s keys from the OS (once per thread, then
28        // varies them per instance), so hashing nothing gives a fresh seed.
29        Self::new(RandomState::new().build_hasher().finish())
30    }
31}
32
33impl Sampler for RandomSampler {
34    fn indices(&mut self, dataset_len: usize) -> Vec<usize> {
35        let mut out = Vec::with_capacity(dataset_len);
36        for i in 0..dataset_len {
37            // Inside-out Fisher-Yates: pick j in [0, i].
38            // If j == i: push i.  Otherwise: copy out[j] to end, write i to out[j].
39            let j = self.rng.usize(0..=i);
40            if j == out.len() {
41                out.push(i);
42            } else {
43                out.push(out[j]);
44                out[j] = i;
45            }
46        }
47        out
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use std::collections::HashSet;
54
55    use crate::sampler::{RandomSampler, Sampler};
56
57    #[test]
58    fn random_same_seed_same_order() {
59        let mut a = RandomSampler::new(42);
60        let mut b = RandomSampler::new(42);
61        assert_eq!(a.indices(20), b.indices(20));
62    }
63
64    #[test]
65    fn random_different_seeds_different_order() {
66        let mut a = RandomSampler::new(1);
67        let mut b = RandomSampler::new(2);
68        assert_ne!(a.indices(20), b.indices(20));
69    }
70
71    #[test]
72    fn random_is_permutation() {
73        let mut s = RandomSampler::new(0);
74        let n = 50;
75        let indices = s.indices(n);
76        assert_eq!(indices.len(), n);
77        let set: HashSet<usize> = indices.into_iter().collect();
78        assert_eq!(set.len(), n);
79        assert!(set.iter().all(|&i| i < n));
80    }
81
82    #[test]
83    fn random_advances_rng_across_epochs() {
84        let mut s = RandomSampler::new(7);
85        let epoch1 = s.indices(30);
86        let epoch2 = s.indices(30);
87        assert_ne!(epoch1, epoch2);
88    }
89
90    #[test]
91    fn random_from_entropy_differs_per_sampler() {
92        let orders: Vec<Vec<usize>> = (0..4)
93            .map(|_| RandomSampler::from_entropy().indices(64))
94            .collect();
95        assert!(orders.windows(2).all(|w| w[0] != w[1]));
96    }
97
98    #[test]
99    fn random_empty_dataset() {
100        let mut s = RandomSampler::new(0);
101        assert_eq!(s.indices(0), Vec::<usize>::new());
102    }
103}
dataloader_rs/sampler/random.rs

dataloader_rs/sampler/
random.rs
