burn-onnx 0.21.0-pre.3

Library for importing ONNX models 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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

// Individual node implementations
// Module declarations (unfortunately must be manual due to Rust's module system)

// Prelude module with common imports for node implementations
pub(crate) mod prelude;

pub(crate) mod broadcast_helpers;
mod subgraph_helper;

// Test helpers for node code generation tests
#[cfg(test)]
pub(crate) mod test_helpers;

pub(crate) mod abs;
pub(crate) mod acos;
pub(crate) mod acosh;
pub(crate) mod add;
pub(crate) mod argmax;
pub(crate) mod argmin;
pub(crate) mod asin;
pub(crate) mod asinh;
pub(crate) mod atan;
pub(crate) mod atanh;
pub(crate) mod attention;
pub(crate) mod avg_pool1d;
pub(crate) mod avg_pool2d;
pub(crate) mod batch_norm;
pub(crate) mod bernoulli;
pub(crate) mod bitshift;
pub(crate) mod bitwiseand;
pub(crate) mod bitwisenot;
pub(crate) mod bitwiseor;
pub(crate) mod bitwisexor;
pub(crate) mod bool_and;
pub(crate) mod bool_or;
pub(crate) mod bool_xor;
pub(crate) mod cast;
pub(crate) mod cast_like;
pub(crate) mod ceil;
pub(crate) mod celu;
pub(crate) mod clip;
pub(crate) mod col2im;
pub(crate) mod concat;
pub(crate) mod constant;
pub(crate) mod constant_of_shape;
pub(crate) mod conv1d;
pub(crate) mod conv2d;
pub(crate) mod conv3d;
pub(crate) mod conv_transpose_1d;
pub(crate) mod conv_transpose_2d;
pub(crate) mod conv_transpose_3d;
pub(crate) mod cos;
pub(crate) mod cosh;
pub(crate) mod cumsum;
pub(crate) mod deform_conv;
pub(crate) mod depth_to_space;
pub(crate) mod dequantize_linear;
pub(crate) mod det;
pub(crate) mod div;
pub(crate) mod dropout;
pub(crate) mod einsum;
pub(crate) mod elu;
pub(crate) mod equal;
pub(crate) mod erf;
pub(crate) mod exp;
pub(crate) mod expand;
pub(crate) mod eye_like;
pub(crate) mod flatten;
pub(crate) mod floor;
pub(crate) mod gather;
pub(crate) mod gather_elements;
pub(crate) mod gathernd;
pub(crate) mod gelu;
pub(crate) mod gemm;
pub(crate) mod global_avg_pool;
pub(crate) mod greater;
pub(crate) mod greater_equal;
pub(crate) mod grid_sample;
pub(crate) mod group_norm;
pub(crate) mod gru;
pub(crate) mod hard_sigmoid;
pub(crate) mod hard_swish;
pub(crate) mod hardmax;
pub(crate) mod identity;
pub(crate) mod if_node;
pub(crate) mod instance_norm;
pub(crate) mod is_inf;
pub(crate) mod is_nan;
pub(crate) mod layer_norm;
pub(crate) mod leaky_relu;
pub(crate) mod linear;
pub(crate) mod log;
pub(crate) mod log_softmax;
pub(crate) mod loop_node;
pub(crate) mod lower;
pub(crate) mod lower_equal;
pub(crate) mod lp_pool1d;
pub(crate) mod lp_pool2d;
pub(crate) mod lrn;
pub(crate) mod lstm;
pub(crate) mod matmul;
pub(crate) mod matmul_integer;
pub(crate) mod max;
pub(crate) mod max_pool1d;
pub(crate) mod max_pool2d;
pub(crate) mod mean;
pub(crate) mod min;
pub(crate) mod mish;
pub(crate) mod modulo;
pub(crate) mod mul;
pub(crate) mod neg;
pub(crate) mod nonzero;
pub(crate) mod not;
pub(crate) mod one_hot;
pub(crate) mod pad;
pub(crate) mod pow;
pub(crate) mod prelu;
pub(crate) mod quantize_linear;
pub(crate) mod random_normal;
pub(crate) mod random_normal_like;
pub(crate) mod random_uniform;
pub(crate) mod random_uniform_like;
pub(crate) mod range;
pub(crate) mod reciprocal;
pub(crate) mod reduce;
pub(crate) mod relu;
pub(crate) mod reshape;
pub(crate) mod resize;
pub(crate) mod rnn;
pub(crate) mod round;
pub(crate) mod scan_node;
pub(crate) mod scatter_elements;
pub(crate) mod scatter_nd;
pub(crate) mod selu;
pub(crate) mod shape;
pub(crate) mod shrink;
pub(crate) mod sigmoid;
pub(crate) mod sign;
pub(crate) mod sin;
pub(crate) mod sinh;
pub(crate) mod size;
pub(crate) mod slice;
pub(crate) mod softmax;
pub(crate) mod softplus;
pub(crate) mod softsign;
pub(crate) mod space_to_depth;
pub(crate) mod split;
pub(crate) mod sqrt;
pub(crate) mod squeeze;
pub(crate) mod sub;
pub(crate) mod sum;
pub(crate) mod swish;
pub(crate) mod tan;
pub(crate) mod tanh;
pub(crate) mod thresholded_relu;
pub(crate) mod tile;
pub(crate) mod top_k;
pub(crate) mod transpose;
pub(crate) mod trilu;
pub(crate) mod unsqueeze;
pub(crate) mod where_op;

// Node codegen traits and utilities are in parent module
// Re-export them here for convenience
pub(crate) use super::node_traits::{NodeCodegen, arg_to_ident};