[TF FE] Support EmptyTensorList and TensorListPushBack operations (#16183)

* [TF FE] Support EmptyTensorList and TensorListPushBack operations

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

* Rename a script to generate the test model

* Correct test model generating script

---------

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

src/frontends/tensorflow/src/op_table.cpp

@@ -133,6 +133,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
         {"DynamicPartition", translate_dynamic_partition_op},
         {"Einsum", translate_einsum_op},
         {"Elu", translate_elu_op},
+        {"EmptyTensorList", translate_tensor_list_reserve_op},
         {"ExpandDims", translate_expand_dims_op},
         {"ExtractImagePatches", translate_extract_image_patches_op},
         {"FakeQuantWithMinMaxVars", translate_fake_quant_op},
@@ -238,6 +239,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
         {"StridedSlice", translate_strided_slice_op},
         {"TensorListFromTensor", translate_tensor_list_from_tensor_op},
         {"TensorListGetItem", translate_tensor_list_get_item_op},
+        {"TensorListPushBack", translate_tensor_list_push_back_op},
         {"TensorListSetItem", translate_tensor_list_set_item_op},
         {"TensorListStack", translate_tensor_list_stack_op},
         {"TensorListReserve", translate_tensor_list_reserve_op},

src/frontends/tensorflow/tests/convert_tricky_models.cpp

@@ -349,4 +349,20 @@ TEST_F(TransformationTestsF, ModelWithMultioutputBodyGraphNode) {
         auto add = make_shared<Add>(topk->output(1), const_ten);
         model_ref = make_shared<Model>(OutputVector{add}, ParameterVector{x, y});
     }
-}
+}
+
+TEST_F(TransformationTestsF, ModelWithEmptyTensorListAndPushBack) {
+    { model = convert_model("empty_tensor_list/empty_tensor_list.pb"); }
+    {
+        auto x = make_shared<Parameter>(f32, Shape{2, 3, 5});
+        auto minus_one_const = make_shared<Constant>(i32, Shape{1}, -1);
+        auto x_flatten = make_shared<Reshape>(x, minus_one_const, false);
+        auto zero_const = make_shared<Constant>(i32, Shape{1}, 0);
+        auto x_unsqueeze_flatten = make_shared<Unsqueeze>(x_flatten, zero_const);
+        auto empty_const = make_shared<Constant>(f32, Shape{0, 30}, vector<float>{});
+        auto list_push_back = make_shared<Concat>(OutputVector{empty_const, x_unsqueeze_flatten}, 0);
+        auto recover_item_shape = make_shared<Constant>(i32, Shape{4}, vector<int32_t>{1, 2, 3, 5});
+        auto recover_item = make_shared<Reshape>(list_push_back, recover_item_shape, false);
+        model_ref = make_shared<Model>(OutputVector{recover_item}, ParameterVector{x});
+    }
+}

src/frontends/tensorflow/tests/test_models/gen_scripts/generate_empty_tensor_list.py

@@ -0,0 +1,25 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+
+import tensorflow as tf
+
+tf.compat.v1.reset_default_graph()
+
+# Create the graph and model
+with tf.compat.v1.Session() as sess:
+    tensor = tf.compat.v1.placeholder(tf.float32, [2, 3, 5], name='data')
+    element_shape = tf.constant([2, 3, 5], dtype=tf.int32)
+    max_num_elements = tf.constant(10, dtype=tf.int32)
+    empty_tensor_list = tf.raw_ops.EmptyTensorList(element_shape=element_shape, max_num_elements=max_num_elements,
+                                                   element_dtype=tf.float32)
+    tensor_list_push_back = tf.raw_ops.TensorListPushBack(input_handle=empty_tensor_list, tensor=tensor)
+    tensor_list_stack = tf.raw_ops.TensorListStack(input_handle=tensor_list_push_back, element_shape=element_shape,
+                                                   element_dtype=tf.float32)
+
+    tf.compat.v1.global_variables_initializer()
+    tf_net = sess.graph_def
+
+tf.io.write_graph(tf_net, os.path.join(sys.argv[1], "empty_tensor_list"), 'empty_tensor_list.pb', False)

src/frontends/tensorflow_common/include/common_op_table.hpp

@@ -129,6 +129,7 @@ OP_CONVERTER(translate_strided_slice_op);
 OP_CONVERTER(translate_sqrt_op);
 OP_CONVERTER(translate_tensor_list_from_tensor_op);
 OP_CONVERTER(translate_tensor_list_get_item_op);
+OP_CONVERTER(translate_tensor_list_push_back_op);
 OP_CONVERTER(translate_tensor_list_reserve_op);
 OP_CONVERTER(translate_tensor_list_set_item_op);
 OP_CONVERTER(translate_tensor_list_stack_op);

src/frontends/tensorflow_common/src/op/tensor_list_operations.cpp

@@ -14,7 +14,7 @@ namespace frontend {
 namespace tensorflow {
 namespace op {
 OutputVector translate_tensor_list_reserve_op(const NodeContext& node) {
-    default_op_checks(node, 2, {"TensorListReserve"});
+    default_op_checks(node, 2, {"TensorListReserve", "EmptyTensorList"});
     auto element_dtype = node.get_attribute<element::Type>("element_dtype");
 
     // always reserve an empty constant of rank equal to two
@@ -122,6 +122,38 @@ OutputVector translate_tensor_list_set_item_op(const NodeContext& node) {
     return {scatter_update};
 }
 
+OutputVector translate_tensor_list_push_back_op(const NodeContext& node) {
+    default_op_checks(node, 2, {"TensorListPushBack"});
+    auto input_handle = node.get_input(0);
+    auto tensor = node.get_input(1);
+
+    // flatten item to be inserted since
+    // the tensor list saves elements in the flatten form
+    // because we want to cover a case of dynamic rank tensor list
+    // the real shape of the tensor elements will be restored by TensorListStack operations
+    auto new_tensor_shape = make_shared<Constant>(element::i32, Shape{1}, -1);
+    tensor = make_shared<Reshape>(tensor, new_tensor_shape, false);
+    auto tensor_shape = make_shared<ShapeOf>(tensor, element::i32);
+
+    // reshape the tensor list to the shape [num_elements, -1]
+    // that is because in the first iteration we have empty constant of a shape [0,0]
+    Output<Node> num_elements = make_shared<ShapeOf>(input_handle, element::i32);
+    auto zero_const = make_shared<Constant>(element::i32, Shape{1}, 0);
+    auto one_const = make_shared<Constant>(element::i32, Shape{1}, 1);
+    num_elements = make_shared<Slice>(num_elements, zero_const, one_const, one_const);
+    auto new_input_handle_shape = make_shared<Concat>(OutputVector{num_elements, tensor_shape}, 0);
+    input_handle = make_shared<Reshape>(input_handle, new_input_handle_shape, false);
+
+    // unsqueeze tensor to be inserted into the list
+    tensor = make_shared<Unsqueeze>(tensor, zero_const);
+
+    // insert the tensor into the end
+    auto updated_list = make_shared<Concat>(OutputVector{input_handle, tensor}, 0);
+
+    set_node_name(node.get_name(), updated_list);
+    return {updated_list};
+}
+
 }  // namespace op
 }  // namespace tensorflow
 }  // namespace frontend