[exo] type and shape inference for Relu6 (#8171)

This adds type and shape inference for Relu6.

Signed-off-by: Hyun Sik Yoon <hyunsik.yoon@samsung.com>

diff --git a/compiler/exo/src/Dialect/Service/TFLShapeInferenceRule.cpp b/compiler/exo/src/Dialect/Service/TFLShapeInferenceRule.cpp
index c4c155a..b42bc37 100644 (file)
--- a/compiler/exo/src/Dialect/Service/TFLShapeInferenceRule.cpp
+++ b/compiler/exo/src/Dialect/Service/TFLShapeInferenceRule.cpp
@@ -305,7 +305,12 @@ public:
     return loco::NodeShape{input_shape};
   }
 
-  // TODO TFLRelu6
+  loco::NodeShape visit(const locoex::TFLRelu6 *node) final
+  {
+    auto input_shape = loco::shape_get(node->features()).as<loco::TensorShape>();
+
+    return loco::NodeShape{input_shape};
+  }
 
   // TODO TFLReshape
 

diff --git a/compiler/exo/src/Dialect/Service/TFLTypeInferenceRule.cpp b/compiler/exo/src/Dialect/Service/TFLTypeInferenceRule.cpp
index cde801f..31d96be 100644 (file)
--- a/compiler/exo/src/Dialect/Service/TFLTypeInferenceRule.cpp
+++ b/compiler/exo/src/Dialect/Service/TFLTypeInferenceRule.cpp
@@ -72,7 +72,10 @@ struct TypeInferenceAlgorithm final : public locoex::TFLNodeVisitor<loco::DataTy
     return loco::dtype_get(node->features());
   }
 
-  // TODO TFLRelu6
+  loco::DataType visit(const locoex::TFLRelu6 *node) final
+  {
+    return loco::dtype_get(node->features());
+  }
 
   // TODO TFLReshape