[TF FE] Post leftovers to support the MUSE model in SavedModel format (#16520)

* [TF FE] Post leftovers to support the MUSE model in SavedModel format

It contains tests imitating a case with Tokenizer extension and raised problems:
setting custom type for body graph Parameter, named ports for RaggedTensorToSparse
and Unique operations.

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

* Update src/frontends/tensorflow/tests/convert_tricky_models.cpp

---------

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

src/frontends/tensorflow/src/decoder_proto.cpp

@@ -306,6 +306,11 @@ void parse_producer_name(const std::string& producer_port_name,
         {"FusedBatchNormV2:batch_variance", 2},
         {"FusedBatchNormV3:batch_mean", 1},
         {"FusedBatchNormV3:batch_variance", 2},
+        {"Unique:y", 0},
+        {"Unique:idx", 1},
+        {"RaggedTensorToSparse:sparse_indices", 0},
+        {"RaggedTensorToSparse:sparse_values", 1},
+        {"RaggedTensorToSparse:sparse_dense_shape", 2},
     };
     // Body graph nodes may have two colons `:` input names, for example,
     // `TopKV2Name:indices:0` means that producer operation name is `TopKV2Name`

src/frontends/tensorflow/src/translate_session.cpp

@@ -66,7 +66,16 @@ void TranslateSession::inject_body_model(std::shared_ptr<ov::Model> body_model,
                             "inputs and arguments to the function " +
                                 operation_type + " do not match.");
     for (size_t param_ind = 0; param_ind < body_parameters.size(); ++param_ind) {
+        auto orig_type = body_parameters[param_ind]->get_element_type();
         body_parameters[param_ind]->output(0).replace(ov_inputs[param_ind]);
+        if (auto ext_parameter = as_type_ptr<ov::opset8::Parameter>(ov_inputs[param_ind].get_node_shared_ptr())) {
+            // save type of a Parameter as converted in the body
+            // this is important if the external conversion extension is applied to body graph node
+            // with setting its own type
+            if (orig_type != element::dynamic) {
+                ext_parameter->set_element_type(orig_type);
+            }
+        }
     }
     for (const auto& result_node : body_model->get_results()) {
         ov_outputs.push_back(result_node->input_value(0));

src/frontends/tensorflow/tests/convert_tricky_models.cpp

@@ -3,6 +3,7 @@
 //
 
 #include <openvino/frontend/exception.hpp>
+#include <openvino/frontend/extension.hpp>
 #include <openvino/frontend/manager.hpp>
 #include <openvino/opsets/opset10.hpp>
 #include <transformations/common_optimizations/moc_transformations.hpp>
@@ -20,12 +21,15 @@ using namespace ov::opset10;
 using namespace ov::frontend;
 
 namespace {
-shared_ptr<Model> convert_model(const string& model_path) {
+shared_ptr<Model> convert_model(const string& model_path, const ConversionExtension::Ptr& conv_ext = nullptr) {
     FrontEndManager fem;
     auto front_end = fem.load_by_framework(TF_FE);
     if (!front_end) {
         throw "TensorFlow Frontend is not initialized";
     }
+    if (conv_ext) {
+        front_end->add_extension(conv_ext);
+    }
     auto model_filename = FrontEndTestUtils::make_model_path(string(TEST_TENSORFLOW_MODELS_DIRNAME) + model_path);
     auto input_model = front_end->load(model_filename);
     if (!input_model) {
@@ -38,6 +42,36 @@ shared_ptr<Model> convert_model(const string& model_path) {
 
     return model;
 }
+
+ov::OutputVector fake_translator_ragged_tensor_to_sparse(const ov::frontend::NodeContext& node) {
+    // NOTE: pay attention that this is a fake translator for RaggedTensorToSparse
+    // only serves for testing purposes
+    FRONT_END_GENERAL_CHECK(node.get_input_size() > 1, "RaggedTensorToSparse expects at least two inputs.");
+    auto node_name = node.get_name();
+    auto row_splits = node.get_input(0);
+    auto strings = node.get_input(1);
+
+    // Override type of input tensor if this is a Parameter
+    if (auto parameter = as_type_ptr<ov::opset10::Parameter>(strings.get_node_shared_ptr())) {
+        parameter->set_partial_shape(ov::PartialShape{Dimension()});
+        parameter->set_element_type(ov::element::u8);
+        parameter->validate_and_infer_types();
+    }
+
+    row_splits = make_shared<ConvertLike>(row_splits, strings);
+    auto const_one = make_shared<Constant>(row_splits.get_element_type(), Shape{}, 1);
+    Output<Node> mul = make_shared<Multiply>(row_splits, const_one);
+    auto const_two = make_shared<Constant>(ov::element::u8, Shape{}, 2);
+    Output<Node> add = make_shared<Add>(strings, const_two);
+    auto const_three = make_shared<Constant>(ov::element::u8, Shape{}, 3);
+    Output<Node> sub = make_shared<Subtract>(strings, const_three);
+
+    mul.get_tensor().add_names({node_name + ":0"});
+    add.get_tensor().add_names({node_name + ":1"});
+    sub.get_tensor().add_names({node_name + ":2"});
+
+    return {mul, add, sub};
+}
 }  // namespace
 
 TEST(FrontEndConvertTrickyModels, undefined_input_shape) {
@@ -387,3 +421,47 @@ TEST_F(TransformationTestsF, ModelWithAssertNode) {
         model_ref = make_shared<Model>(OutputVector{add}, ParameterVector{x, y});
     }
 }
+
+TEST_F(TransformationTestsF, PartitionedCallWithUnique) {
+    // This test aims to test named output ports for Unique operation
+    { model = convert_model("partitioned_call_with_unique/partitioned_call_with_unique.pb"); }
+    {
+        auto x = make_shared<Parameter>(f32, Shape{5});
+        auto relu = make_shared<Relu>(x);
+        auto unique = make_shared<Unique>(relu, false, i32);
+        auto const_one = make_shared<Constant>(i32, Shape{}, 1);
+        auto add = make_shared<Add>(unique->output(2), const_one);
+        auto sigmoid = make_shared<Sigmoid>(unique->output(0));
+        model_ref = make_shared<Model>(OutputVector{sigmoid, add}, ParameterVector{x});
+    }
+}
+
+TEST_F(TransformationTestsF, RaggedTensorToSparse) {
+    // This test aims to test named output ports for RaggedTensorToSparse operation
+    // also, it tests propagation of custom type (specified in the extension) to Parameter node in the parent graph
+    {
+        // create FAKE conversion extension for RaggedTensorToSparse
+        auto conv_ext = std::make_shared<ov::frontend::ConversionExtension>("RaggedTensorToSparse",
+                                                                            fake_translator_ragged_tensor_to_sparse);
+        model = convert_model("ragged_tensor_to_sparse/ragged_tensor_to_sparse.pb", conv_ext);
+    }
+    {
+        auto strings = make_shared<Parameter>(u8, PartialShape{3});
+        auto row_splits = make_shared<Parameter>(i32, PartialShape{5});
+        auto convert_like = make_shared<ConvertLike>(row_splits, strings);
+
+        auto const_one = make_shared<Constant>(u8, Shape{}, 1);
+        Output<Node> mul = make_shared<Multiply>(convert_like, const_one);
+        auto const_three = make_shared<Constant>(u8, Shape{}, 3);
+        Output<Node> sub = make_shared<Subtract>(strings, const_three);
+
+        auto target_shape1 = make_shared<Constant>(i32, Shape{1}, -1);
+        auto reshape1 = make_shared<Reshape>(mul, target_shape1, false);
+        auto target_shape2 = make_shared<Constant>(i32, Shape{1}, -1);
+        auto reshape2 = make_shared<Reshape>(sub, target_shape2, false);
+
+        auto concat = make_shared<Concat>(OutputVector{reshape1, reshape2}, 0);
+
+        model_ref = make_shared<Model>(OutputVector{concat}, ParameterVector{row_splits, strings});
+    }
+}

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

@@ -0,0 +1,33 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+
+import tensorflow.compat.v1 as tf
+
+
+def main():
+    tf.compat.v1.reset_default_graph()
+
+    @tf.function
+    def second_func(x):
+        relu = tf.raw_ops.Relu(features=x)
+        y, idx = tf.raw_ops.Unique(x=relu)
+        return y, idx
+
+    @tf.function
+    def first_func(x):
+        y, idx = second_func(x)
+        sigmoid = tf.raw_ops.Sigmoid(x=y)
+        const_one = tf.constant(1, dtype=tf.int32)
+        add = tf.raw_ops.AddV2(x=idx, y=const_one)
+        return sigmoid, add
+
+    tf_net = first_func.get_concrete_function(tf.constant([1, 2, 3, 4, 5], dtype=tf.float32)).graph.as_graph_def()
+    tf.io.write_graph(tf_net, os.path.join(sys.argv[1], "partitioned_call_with_unique"),
+                      "partitioned_call_with_unique.pb", False)
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,35 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+
+import tensorflow.compat.v1 as tf
+
+
+def main():
+    tf.compat.v1.reset_default_graph()
+
+    @tf.function
+    def second_func(strings, row_splits):
+        sparse_indices, sparse_values, sparse_dense_shape = tf.raw_ops.RaggedTensorToSparse(
+            rt_nested_splits=[row_splits],
+            rt_dense_values=strings)
+        return sparse_indices, sparse_values, sparse_dense_shape
+
+    @tf.function
+    def first_func(strings, row_splits):
+        sparse_indices, _, sparse_dense_shape = second_func(strings, row_splits)
+        sparse_indices = tf.raw_ops.Reshape(tensor=sparse_indices, shape=[-1])
+        sparse_dense_shape = tf.raw_ops.Reshape(tensor=sparse_dense_shape, shape=[-1])
+        concat = tf.raw_ops.Concat(concat_dim=0, values=[sparse_indices, sparse_dense_shape])
+        return concat
+
+    tf_net = first_func.get_concrete_function(tf.constant(["abc", "bcd", "cc"], dtype=tf.string),
+                                              tf.constant([0, 2, 2, 3, 3], dtype=tf.int32)).graph.as_graph_def()
+    tf.io.write_graph(tf_net, os.path.join(sys.argv[1], "ragged_tensor_to_sparse"),
+                      "ragged_tensor_to_sparse.pb", False)
+
+
+if __name__ == "__main__":
+    main()