burn-cubecl 0.22.0-pre.4

Generic backend that can be compiled just-in-time to any shader language target
Documentation 
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

use crate::{CubeDevice, ops::numeric::empty_device_dtype, tensor::CubeTensor};
use burn_backend::cubecl::dtype_to_storage_type;
use burn_backend::{DType, Shape, TensorMetadata};

/// Pseudo-random generator with uniform distribution
pub fn random_uniform(
    shape: Shape,
    device: &CubeDevice,
    lower_bound: f32,
    upper_bound: f32,
    dtype: DType,
) -> CubeTensor {
    let client = device.client();
    let output = empty_device_dtype(client.clone(), device.clone(), shape, dtype);

    cubek::random::random_uniform(
        &client,
        lower_bound,
        upper_bound,
        output.clone().binding(),
        dtype_to_storage_type(dtype),
    )
    .expect("Kernel to never fail");

    output
}

/// Pseudo-random generator for uniform distribution, based on
/// another tensor.
pub fn random_like_uniform(
    tensor: &CubeTensor,
    lower_bound: f32,
    upper_bound: f32,
    dtype: DType,
) -> CubeTensor {
    random_uniform(
        tensor.shape(),
        &tensor.device,
        lower_bound,
        upper_bound,
        dtype,
    )
}