From b44f915a9d3e6c013e0519f13f5c55d8c2e97df6 Mon Sep 17 00:00:00 2001
From: Roman Kazantsev <roman.kazantsev@intel.com>
Date: Wed, 2 Aug 2023 17:28:39 +0400
Subject: [PATCH] [TF FE] Support ShapeN operation (#18913)

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
---
 src/frontends/tensorflow/src/op_table.cpp     |  1 +
 .../tensorflow_common/src/op/shape.cpp        | 33 ++++++++++++----
 .../tensorflow_tests/test_tf_ShapeN.py        | 38 +++++++++++++++++++
 3 files changed, 64 insertions(+), 8 deletions(-)
 create mode 100644 tests/layer_tests/tensorflow_tests/test_tf_ShapeN.py

diff --git a/src/frontends/tensorflow/src/op_table.cpp b/src/frontends/tensorflow/src/op_table.cpp
index 7ec883a2f39..6ab15f47ea4 100644
--- a/src/frontends/tensorflow/src/op_table.cpp
+++ b/src/frontends/tensorflow/src/op_table.cpp
@@ -233,6 +233,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
         {"Select", CreatorFunction(translate_select_op)},
         {"SelectV2", CreatorFunction(translate_select_v2_op)},
         {"Shape", CreatorFunction(translate_shape_op)},
+        {"ShapeN", CreatorFunction(translate_shape_op)},
         {"Size", CreatorFunction(translate_size_op)},
         {"Slice", CreatorFunction(translate_slice_op)},
         {"Snapshot", CreatorFunction(translate_identity_op)},
diff --git a/src/frontends/tensorflow_common/src/op/shape.cpp b/src/frontends/tensorflow_common/src/op/shape.cpp
index 3aa5d762cdb..749009c67ea 100644
--- a/src/frontends/tensorflow_common/src/op/shape.cpp
+++ b/src/frontends/tensorflow_common/src/op/shape.cpp
@@ -3,24 +3,41 @@
 //
 
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/shape_of.hpp"
 
 using namespace std;
 using namespace ov;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
-ov::OutputVector translate_shape_op(const NodeContext& node) {
-    default_op_checks(node, 1, {"Shape", "SHAPE"});
-    auto input = node.get_input(0);
+OutputVector translate_shape_op(const NodeContext& node) {
+    default_op_checks(node, 1, {"Shape", "ShapeN", "SHAPE"});
+    auto input_size = static_cast<int>(node.get_input_size());
     auto out_type = node.get_attribute<element::Type>("out_type", element::i32);
-    auto shapeof = make_shared<ShapeOf>(input, out_type);
-    set_node_name(node.get_name(), shapeof);
-    return {shapeof};
+    auto node_name = node.get_name();
+
+    if (input_size == 1) {
+        auto input = node.get_input(0);
+        auto shapeof = make_shared<v3::ShapeOf>(input, out_type);
+        set_node_name(node_name, shapeof);
+        return {shapeof};
+    }
+
+    OutputVector outputs;
+    for (int input_ind = 0; input_ind < input_size; ++input_ind) {
+        auto input = node.get_input(input_ind);
+        auto shapeof = make_shared<v3::ShapeOf>(input, out_type);
+        shapeof->set_friendly_name(node_name + "_" + to_string(input_ind));
+        auto shapeof_output = shapeof->output(0);
+        set_out_name({node_name + ":" + to_string(input_ind)}, shapeof_output);
+        outputs.push_back(shapeof_output);
+    }
+
+    return outputs;
 }
 
 }  // namespace op
diff --git a/tests/layer_tests/tensorflow_tests/test_tf_ShapeN.py b/tests/layer_tests/tensorflow_tests/test_tf_ShapeN.py
new file mode 100644
index 00000000000..dd358945ea2
--- /dev/null
+++ b/tests/layer_tests/tensorflow_tests/test_tf_ShapeN.py
@@ -0,0 +1,38 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import tensorflow as tf
+from common.tf_layer_test_class import CommonTFLayerTest
+
+
+class TestShapeN(CommonTFLayerTest):
+    def create_shape_n_net(self, input_shapes, out_type):
+        tf.compat.v1.reset_default_graph()
+        # Create the graph and model
+        with tf.compat.v1.Session() as sess:
+            inputs = []
+            for ind, input_shape in enumerate(input_shapes):
+                inputs.append(tf.compat.v1.placeholder(tf.float32, input_shape, 'input_{}'.format(ind)))
+
+            shapen = tf.raw_ops.ShapeN(input=inputs, out_type=out_type)
+            tf.raw_ops.ConcatV2(values=shapen, axis=0)
+
+            tf.compat.v1.global_variables_initializer()
+            tf_net = sess.graph_def
+
+        return tf_net, None
+
+    test_data_basic = [
+        dict(input_shapes=[[2, 3], [1]], out_type=tf.int32),
+        dict(input_shapes=[[3], [3, 2, 1], [], [4, 3, 1, 1]], out_type=tf.int64),
+    ]
+
+    @pytest.mark.parametrize("params", test_data_basic)
+    @pytest.mark.precommit_tf_fe
+    @pytest.mark.nightly
+    def test_shape_n_basic(self, params, ie_device, precision, ir_version, temp_dir,
+                           use_new_frontend, use_old_api):
+        self._test(*self.create_shape_n_net(**params),
+                   ie_device, precision, ir_version, temp_dir=temp_dir,
+                   use_new_frontend=use_new_frontend, use_old_api=use_old_api)