cubek-random 0.3.0

CubeK: Random Number Generation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

use cubecl::{prelude::*, std::tensor::TensorHandle};
use cubek_random::*;

#[test]
fn empirical_mean_close_to_expectation() {
    let shape = &[100, 100];
    let mean = 10.;
    let std = 2.;

    let output_data = get_random_normal_data(shape, mean, std);

    assert_mean_approx_equal(&output_data, mean);

    let shape = &[1000, 1000];
    let mean = 0.;
    let std = 1.;

    let output_data = get_random_normal_data(shape, mean, std);

    assert_mean_approx_equal(&output_data, mean);
}

#[test]
fn normal_respects_68_95_99_rule() {
    // https://en.wikipedia.org/wiki/68%E2%80%9395%E2%80%9399.7_rule
    let shape = &[1000, 1000];
    let mu = 0.;
    let s = 1.;

    let output_data = get_random_normal_data(shape, mu, s);

    assert_normal_respects_68_95_99_rule(&output_data, mu, s);
}

fn get_random_normal_data(shape: &[usize], mean: f32, std: f32) -> Vec<TestDType> {
    let client = cubecl::test_device().client();
    let output = TensorHandle::empty(&client, shape.to_vec(), TestDType::elem_type_native());

    with_seed(0, || {
        random_normal(
            &client,
            mean,
            std,
            output.clone().binding(),
            TestDType::elem_type_native(),
        )
    })
    .unwrap();

    let output_data = client.read_one_unchecked_tensor(output.into_copy_descriptor());
    let output_data = TestDType::from_bytes(&output_data);

    output_data.to_owned()
}