refactor: burn-import unary operators (#548)

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

@@ -1,8 +1,7 @@
 use super::{
     add::AddNode, batch_norm::BatchNormNode, constant::ConstantNode, conv2d::Conv2dNode,
-    equal::EqualNode, flatten::FlattenNode, linear::LinearNode, log_softmax::LogSoftmaxNode,
-    matmul::MatmulNode, max_pool2d::MaxPool2dNode, relu::ReLUNode, reshape::ReshapeNode,
-    sigmoid::SigmoidNode,
+    equal::EqualNode, linear::LinearNode, matmul::MatmulNode, max_pool2d::MaxPool2dNode,
+    reshape::ReshapeNode, unary::UnaryNode,
 };
 use crate::burn::{BurnImports, Scope, Type};
 use burn::record::PrecisionSettings;
@@ -78,13 +77,10 @@ pub enum Node<PS: PrecisionSettings> {
     MaxPool2d(MaxPool2dNode),
     Linear(LinearNode<PS>),
     BatchNorm(BatchNormNode<PS>),
-    ReLU(ReLUNode),
-    Flatten(FlattenNode),
-    LogSoftmax(LogSoftmaxNode),
     Constant(ConstantNode),
     Equal(EqualNode),
+    Unary(UnaryNode),
     Reshape(ReshapeNode),
-    Sigmoid(SigmoidNode),
 }
 
 macro_rules! match_all {
@@ -96,13 +92,10 @@ macro_rules! match_all {
             Node::MaxPool2d(node) => $func(node),
             Node::Linear(node) => $func(node),
             Node::BatchNorm(node) => $func(node),
-            Node::ReLU(node) => $func(node),
-            Node::Flatten(node) => $func(node),
-            Node::LogSoftmax(node) => $func(node),
             Node::Constant(node) => $func(node),
             Node::Equal(node) => $func(node),
             Node::Reshape(node) => $func(node),
-            Node::Sigmoid(node) => $func(node),
+            Node::Unary(node) => $func(node),
         }
     }};
 }
@@ -126,12 +119,9 @@ impl<PS: PrecisionSettings> Node<PS> {
             Node::MaxPool2d(_) => "max_pool2d",
             Node::Linear(_) => "linear",
             Node::BatchNorm(_) => "batch_norm",
-            Node::ReLU(_) => "relu",
-            Node::Flatten(_) => "flatten",
-            Node::LogSoftmax(_) => "log_softmax",
             Node::Equal(_) => "equal",
             Node::Reshape(_) => "reshape",
-            Node::Sigmoid(_) => "sigmoid",
+            Node::Unary(unary) => unary.kind.as_str(),
         }
     }
 }
@@ -264,9 +254,9 @@ pub(crate) mod tests {
                 module::Module,
                 tensor::{backend::Backend, Tensor},
             };
-            use burn::nn::PaddingConfig2d;
-            use burn::nn::conv::Conv2d;
             use burn::nn::conv::Conv2dConfig;
+            use burn::nn::conv::Conv2d;
+            use burn::nn::PaddingConfig2d;
 
             #[derive(Module, Debug)]
             pub struct Model <B: Backend> {

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

@@ -5,14 +5,11 @@ pub(crate) mod batch_norm;
 pub(crate) mod constant;
 pub(crate) mod conv2d;
 pub(crate) mod equal;
-pub(crate) mod flatten;
 pub(crate) mod linear;
-pub(crate) mod log_softmax;
 pub(crate) mod matmul;
 pub(crate) mod max_pool2d;
-pub(crate) mod relu;
 pub(crate) mod reshape;
-pub(crate) mod sigmoid;
+pub(crate) mod unary;
 
 pub(crate) use base::*;