burn-import 0.14.0

Library for importing datamodels into the Burn framework
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

use super::{Node, NodeCodegen};
use crate::burn::{Scope, TensorType, Type};
use burn::record::PrecisionSettings;
use proc_macro2::TokenStream;
use quote::quote;

#[derive(Debug, Clone)]
pub struct RandomNormalNode {
    pub mean: f64,
    pub scale: f64,
    pub output_ty: TensorType,
}

impl RandomNormalNode {
    pub fn new(output_ty: TensorType, mean: f64, scale: f64) -> Self {
        Self {
            mean,
            scale,
            output_ty,
        }
    }

    fn get_output_shape(&self) -> TokenStream {
        let shape_it = self
            .output_ty
            .shape
            .as_ref()
            .expect("RandomNormal output has no shape!")
            .iter();
        quote! { Shape::new([#(#shape_it),*]) }
    }

    fn get_distribution(&self) -> TokenStream {
        let std_deviation = self.scale; // ONNX spec defines `scale` == `standard deviation`
        let mean = self.mean;
        quote! { Distribution::Normal(#mean, #std_deviation) }
    }
}

impl<PS: PrecisionSettings> NodeCodegen<PS> for RandomNormalNode {
    fn input_types(&self) -> Vec<Type> {
        Vec::with_capacity(0)
    }

    fn output_types(&self) -> Vec<Type> {
        vec![Type::Tensor(self.output_ty.clone())]
    }

    fn forward(&self, _scope: &mut Scope, _node_position: usize) -> TokenStream {
        let output = &self.output_ty.name;
        let shape = self.get_output_shape();
        let dist = self.get_distribution();
        quote! {
            let #output = Tensor::random(#shape, #dist, &*self.device);
        }
    }

    fn into_node(self) -> Node<PS> {
        Node::RandomNormal(self)
    }

    fn register_imports(&self, imports: &mut crate::burn::BurnImports) {
        imports.register("burn::tensor::Distribution");
        imports.register("burn::prelude::Shape");
    }
}

#[cfg(test)]
mod tests {
    use burn::record::FullPrecisionSettings;

    use super::*;
    use crate::burn::{
        graph::BurnGraph,
        node::{random_normal::RandomNormalNode, test::assert_tokens},
        TensorKind, TensorType,
    };

    #[test]
    fn test_codegen_nodes() {
        let mut graph = BurnGraph::<FullPrecisionSettings>::default();

        graph.register(RandomNormalNode::new(
            TensorType::new("tensor1", 2, TensorKind::Float, Some(vec![2, 3])),
            0.0f64,
            1.0f64,
        ));

        graph.register_input_output(vec![], vec!["tensor1".to_string()]);

        let expected = quote! {
            use burn::prelude::Shape;
            use burn::tensor::Distribution;
            use burn::{
                module::Module,
                tensor::{backend::Backend, Tensor},
            };

            #[derive(Module, Debug)]
            pub struct Model<B: Backend> {
                phantom: core::marker::PhantomData<B>,
                device: burn::module::Ignored<B::Device>,
            }

            impl<B: Backend> Model <B> {
                #[allow(unused_variables)]
                pub fn new(device: &B::Device) -> Self {
                    Self {
                        phantom: core::marker::PhantomData,
                        device: burn::module::Ignored(device.clone()),
                    }
                }
                #[allow(clippy::let_and_return, clippy::approx_constant)]
                pub fn forward(&self) -> Tensor<B, 2> {
                    let tensor1 = Tensor::random(
                        Shape::new([2usize, 3usize]),
                        Distribution::Normal(0f64, 1f64),
                        &*self.device,
                    );

                    tensor1
                }
            }
        };

        assert_tokens(graph.codegen(), expected);
    }
}