burn-jit 0.16.1

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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

#[burn_tensor_testgen::testgen(reduce)]
mod reduce {
    use super::*;
    use burn_jit::kernel::reduce::{
        reduce, reduce_dim, ArgMax, ArgMin, Mean, Prod, ReduceStrategy, Sum,
    };
    use burn_tensor::{
        backend::Backend, ops::IntTensorOps, Distribution, Int, Shape, Tensor, TensorData,
        TensorPrimitive,
    };

    const RANK: usize = 4;
    const SHAPE: [usize; RANK] = [2, 4, 8, 16];

    #[test]
    fn reduction_argmax_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        for dim in 0..RANK {
            tensor
                .clone()
                .argmax(dim)
                .into_data()
                .assert_eq(&tensor_ref.clone().argmax(dim).into_data(), false);
        }
    }

    #[test]
    fn reduction_argmin_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        for dim in 0..RANK {
            tensor
                .clone()
                .argmin(dim)
                .into_data()
                .assert_eq(&tensor_ref.clone().argmin(dim).into_data(), false);
        }
    }

    #[test]
    fn reduction_mean_dim_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        for dim in 0..RANK {
            tensor
                .clone()
                .mean_dim(dim)
                .into_data()
                .assert_approx_eq_diff(&tensor_ref.clone().mean_dim(dim).into_data(), 1e-6);
        }
    }

    #[test]
    fn reduction_mean_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        tensor
            .clone()
            .mean()
            .into_data()
            .assert_approx_eq_diff(&tensor_ref.clone().mean().into_data(), 1e-6);
    }

    #[test]
    fn reduction_prod_dim_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        for dim in 0..RANK {
            tensor
                .clone()
                .prod_dim(dim)
                .into_data()
                .assert_approx_eq_diff(&tensor_ref.clone().prod_dim(dim).into_data(), 1e-6);
        }
    }

    #[test]
    fn reduction_prod_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        tensor
            .clone()
            .prod()
            .into_data()
            .assert_approx_eq_diff(&tensor_ref.clone().prod().into_data(), 1e-6);
    }

    #[test]
    fn reduction_sum_dim_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        for dim in 0..RANK {
            tensor
                .clone()
                .sum_dim(dim)
                .into_data()
                .assert_approx_eq_diff(&tensor_ref.clone().sum_dim(dim).into_data(), 1e-6);
        }
    }

    #[test]
    fn reduction_sum_should_match_reference_backend() {
        let tensor =
            Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
        let tensor_ref =
            Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
        tensor
            .clone()
            .sum()
            .into_data()
            .assert_approx_eq_diff(&tensor_ref.clone().sum().into_data(), 1e-6);
    }
}