Implement binary operators for burn-import (#532)

burn-import/README.md

@@ -56,12 +56,12 @@ List taken from [here](https://github.com/onnx/onnx/blob/main/docs/Operators.md)
 - [ ] DequantizeLinear
 - [ ] Det
 - [ ] DFT
-- [ ] Div
+- [x] Div
 - [ ] Dropout
 - [ ] DynamicQuantizeLinear
 - [ ] Einsum
 - [ ] Elu
-- [ ] Equal
+- [x] Equal
 - [ ] Erf
 - [ ] Exp
 - [ ] Expand
@@ -189,7 +189,7 @@ List taken from [here](https://github.com/onnx/onnx/blob/main/docs/Operators.md)
 - [ ] Squeeze
 - [ ] STFT
 - [ ] StringNormalizer
-- [ ] Sub
+- [x] Sub
 - [ ] Sum
 - [ ] Tan
 - [ ] Tanh

burn-import/src/burn/node/base.rs

@@ -1,7 +1,7 @@
 use super::{
-    add::AddNode, batch_norm::BatchNormNode, constant::ConstantNode, conv2d::Conv2dNode,
-    equal::EqualNode, linear::LinearNode, matmul::MatmulNode, max_pool2d::MaxPool2dNode,
-    reshape::ReshapeNode, unary::UnaryNode,
+    batch_norm::BatchNormNode, binary::BinaryNode, constant::ConstantNode, conv2d::Conv2dNode,
+    linear::LinearNode, matmul::MatmulNode, max_pool2d::MaxPool2dNode, reshape::ReshapeNode,
+    unary::UnaryNode,
 };
 use crate::burn::{BurnImports, Scope, Type};
 use burn::record::PrecisionSettings;
@@ -71,31 +71,29 @@ pub trait NodeCodegen<PS: PrecisionSettings>: std::fmt::Debug {
 
 #[derive(Debug)]
 pub enum Node<PS: PrecisionSettings> {
-    Add(AddNode),
+    Binary(BinaryNode),
     Matmul(MatmulNode),
     Conv2d(Conv2dNode<PS>),
     MaxPool2d(MaxPool2dNode),
     Linear(LinearNode<PS>),
     BatchNorm(BatchNormNode<PS>),
     Constant(ConstantNode),
-    Equal(EqualNode),
     Unary(UnaryNode),
     Reshape(ReshapeNode),
 }
 
 macro_rules! match_all {
     ($self:expr, $func:expr) => {{
         match $self {
-            Node::Add(node) => $func(node),
             Node::Matmul(node) => $func(node),
             Node::Conv2d(node) => $func(node),
             Node::MaxPool2d(node) => $func(node),
             Node::Linear(node) => $func(node),
             Node::BatchNorm(node) => $func(node),
             Node::Constant(node) => $func(node),
-            Node::Equal(node) => $func(node),
             Node::Reshape(node) => $func(node),
             Node::Unary(node) => $func(node),
+            Node::Binary(node) => $func(node),
         }
     }};
 }
@@ -112,16 +110,15 @@ impl<PS: PrecisionSettings> Serialize for Node<PS> {
 impl<PS: PrecisionSettings> Node<PS> {
     pub fn name(&self) -> &str {
         match self {
-            Node::Add(_) => "add",
             Node::Matmul(_) => "matmul",
             Node::Constant(_) => "constant",
             Node::Conv2d(_) => "conv2d",
             Node::MaxPool2d(_) => "max_pool2d",
             Node::Linear(_) => "linear",
             Node::BatchNorm(_) => "batch_norm",
-            Node::Equal(_) => "equal",
             Node::Reshape(_) => "reshape",
             Node::Unary(unary) => unary.kind.as_str(),
+            Node::Binary(binary) => binary.binary_type.as_str(),
         }
     }
 }
@@ -187,17 +184,13 @@ pub(crate) mod tests {
 
     fn one_node_graph<T: NodeCodegen<FullPrecisionSettings> + 'static>(
         node_gen: T,
-    ) -> BurnGraph<FullPrecisionSettings> {
+        forward: TokenStream,
+    ) {
         let mut graph = BurnGraph::<FullPrecisionSettings>::default();
 
         graph.register(node_gen);
 
-        graph
-    }
-
-    fn unary_operator_expected<const N: usize>(function: TokenStream) -> TokenStream {
-        let tensor_dim = N.to_tokens();
-        quote! {
+        let expected = quote! {
             use burn::{
                 module::Module,
                 tensor::{backend::Backend, Tensor},
@@ -210,25 +203,49 @@ pub(crate) mod tests {
                 pub fn new_with(record: ModelRecord<B>) -> Self {
                     Self { }
                 }
+
                 #[allow(clippy::let_and_return)]
+                #forward
+            }
+        };
+
+        assert_tokens(graph.codegen(), expected);
+    }
+
+    pub(crate) fn codegen_unary_operator<
+        const N: usize,
+        T: NodeCodegen<FullPrecisionSettings> + 'static,
+    >(
+        node_gen: T,
+        function: TokenStream,
+    ) {
+        let forward = |function, tensor_dim| {
+            quote! {
                 pub fn forward(&self, tensor1: Tensor<B, #tensor_dim>) -> Tensor<B, #tensor_dim> {
                     #function
                 }
             }
-        }
+        };
+
+        one_node_graph(node_gen, forward(function, N.to_tokens()));
     }
 
-    pub(crate) fn codegen_unary_operator<
+    pub(crate) fn codegen_binary_operator<
         const N: usize,
         T: NodeCodegen<FullPrecisionSettings> + 'static,
     >(
         node_gen: T,
         function: TokenStream,
     ) {
-        assert_tokens(
-            one_node_graph(node_gen).codegen(),
-            unary_operator_expected::<N>(function),
-        );
+        let forward = |function, tensor_dim| {
+            quote! {
+                pub fn forward(&self, tensor1: Tensor<B, #tensor_dim>, tensor2: Tensor<B, #tensor_dim>) -> Tensor<B, #tensor_dim> {
+                    #function
+                }
+            }
+        };
+
+        one_node_graph(node_gen, forward(function, N.to_tokens()));
     }
 
     #[test]