Skip to main content

burn_cubecl/kernel/prng/
normal.rs

1use crate::{CubeDevice, ops::numeric::empty_device_dtype, tensor::CubeTensor};
2use burn_backend::cubecl::dtype_to_storage_type;
3use burn_backend::{DType, Shape};
4
5/// Pseudo-random generator with uniform distribution
6pub fn random_normal(
7    shape: Shape,
8    device: &CubeDevice,
9    mean: f32,
10    std: f32,
11    dtype: DType,
12) -> CubeTensor {
13    let client = device.client();
14    let output = empty_device_dtype(client.clone(), device.clone(), shape, dtype);
15
16    cubek::random::random_normal(
17        &client,
18        mean,
19        std,
20        output.clone().binding(),
21        dtype_to_storage_type(dtype),
22    )
23    .expect("Kernel to never fail");
24
25    output
26}