[exo-tflite] Shapeinf for EltwiseAdd and EltwiseMul (#6099)

This will add to support ShapeInference for EltwiseAdd and EltwiseMul

Signed-off-by: SaeHie Park <saehie.park@samsung.com>

diff --git a/compiler/exo-tflite/src/ShapeInference.cpp b/compiler/exo-tflite/src/ShapeInference.cpp
index 0109105..2891b55 100644 (file)
--- a/compiler/exo-tflite/src/ShapeInference.cpp
+++ b/compiler/exo-tflite/src/ShapeInference.cpp
@@ -79,6 +79,8 @@ public:
   NODE(BiasEncode)
   NODE(TensorBiasAdd)
   NODE(FeatureBiasAdd)
+  NODE(EltwiseAdd)
+  NODE(EltwiseMul)
 #undef NODE
   // TODO Put all the visit method implementations inside this class declaration
   ShapeDescription visit(loco::ReLU6 *node) { return gd._node_to_shape[node->input()]; }
@@ -425,6 +427,26 @@ ShapeDescription ShapeGetter::visit(loco::FeatureBiasAdd *node)
   return value_shape;
 }
 
+ShapeDescription ShapeGetter::visit(loco::EltwiseAdd *node)
+{
+  const ShapeDescription &lhs_shape = gd._node_to_shape[node->lhs()];
+  const ShapeDescription &rhs_shape = gd._node_to_shape[node->rhs()];
+
+  assert(lhs_shape._dims == rhs_shape._dims);
+
+  return lhs_shape;
+}
+
+ShapeDescription ShapeGetter::visit(loco::EltwiseMul *node)
+{
+  const ShapeDescription &lhs_shape = gd._node_to_shape[node->lhs()];
+  const ShapeDescription &rhs_shape = gd._node_to_shape[node->rhs()];
+
+  assert(lhs_shape._dims == rhs_shape._dims);
+
+  return lhs_shape;
+}
+
 } // namespace
 
 namespace