From 19ff7fba3d30e3658a21ca46fc8a1cd7d2342030 Mon Sep 17 00:00:00 2001
From: Roman Kazantsev <roman.kazantsev@intel.com>
Date: Mon, 21 Aug 2023 14:26:36 +0400
Subject: [PATCH] [TF FE] Fix support of CTCLoss and add tests to pre-commit
 (#19291)

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
---
 .../tensorflow_common/src/op/ctc_loss.cpp     | 73 +++++++++++--------
 .../tensorflow_tests/test_tf_CTCLoss.py       | 30 ++++----
 2 files changed, 55 insertions(+), 48 deletions(-)

diff --git a/src/frontends/tensorflow_common/src/op/ctc_loss.cpp b/src/frontends/tensorflow_common/src/op/ctc_loss.cpp
index b4b72bb407f..1abba8801f2 100644
--- a/src/frontends/tensorflow_common/src/op/ctc_loss.cpp
+++ b/src/frontends/tensorflow_common/src/op/ctc_loss.cpp
@@ -2,14 +2,24 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#include "openvino/op/ctc_loss.hpp"
+
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/broadcast.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/convert.hpp"
+#include "openvino/op/equal.hpp"
+#include "openvino/op/reduce_sum.hpp"
+#include "openvino/op/scatter_nd_update.hpp"
+#include "openvino/op/select.hpp"
+#include "openvino/op/shape_of.hpp"
+#include "openvino/op/slice.hpp"
 
 using namespace std;
 using namespace ov;
-using namespace opset8;
 using namespace ov::frontend;
-using namespace frontend::tensorflow;
+using namespace ov::frontend::tensorflow;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
@@ -17,7 +27,7 @@ namespace tensorflow {
 namespace op {
 
 OutputVector translate_ctc_loss_op(const NodeContext& node) {
-    // This is a translator for CTCLoss v1 aka tf.compat.v1.nn.ctc_loss
+    // this is a translator for CTCLoss v1 aka tf.compat.v1.nn.ctc_loss
     default_op_checks(node, 4, {"CTCLoss"});
     auto logits = node.get_input(0);
     auto decoded_indices = node.get_input(1);
@@ -34,40 +44,41 @@ OutputVector translate_ctc_loss_op(const NodeContext& node) {
         // we need to transpose it into [batch_size, time_size, num_classes] format
         // from [time_size, batch_size, num_classes]
         AxisVector logits_order = {1, 0, 2};
-        logits = tensorflow::make_transpose(logits, logits_order);
+        logits = make_transpose(logits, logits_order);
     }
 
-    // Transform decoded labels from the sparse format into dense format
-    // Convert to the signed type since the mask with minus one is formed below
-    decoded_values = make_shared<Convert>(decoded_values, element::i64);
+    // transform decoded labels from the sparse format into dense format
+    // convert to the signed type since the mask with minus one is formed below
+    decoded_values = make_shared<v0::Convert>(decoded_values, element::i64);
     // OpenVINO ScatterND operation requires indices to be signed
-    decoded_indices = make_shared<Convert>(decoded_indices, element::i64);
+    decoded_indices = make_shared<v0::Convert>(decoded_indices, element::i64);
     // OpenVINO CTCLoss requires logit_length to be signed
-    logit_length = make_shared<Convert>(logit_length, element::i64);
+    logit_length = make_shared<v0::Convert>(logit_length, element::i64);
 
-    auto logits_shape = make_shared<ShapeOf>(logits, element::i64);
-    auto dense_shape = make_shared<Slice>(logits_shape,
-                                          make_shared<Constant>(element::i64, Shape{}, 0),
-                                          make_shared<Constant>(element::i64, Shape{}, 2),
-                                          make_shared<Constant>(element::i64, Shape{}, 1));
-    auto minus_one_value = make_shared<Constant>(element::i64, Shape{}, -1);
-    auto init_decoded_values = make_shared<Broadcast>(minus_one_value, dense_shape);
-    auto decoded_values_dense = make_shared<ScatterNDUpdate>(init_decoded_values, decoded_indices, decoded_values);
+    // compute target labels in a format accepted by OpenVINO CTCLoss
+    auto logits_shape = make_shared<v3::ShapeOf>(logits, element::i64);
+    auto slice_start = make_shared<v0::Constant>(element::i64, Shape{1}, 0);
+    auto slice_end = make_shared<v0::Constant>(element::i64, Shape{1}, 2);
+    auto slice_step = make_shared<v0::Constant>(element::i64, Shape{1}, 1);
+    auto dense_shape = make_shared<v8::Slice>(logits_shape, slice_start, slice_end, slice_step);
+    auto minus_one = make_shared<v0::Constant>(element::i64, Shape{}, -1);
+    auto labels = make_shared<v3::Broadcast>(minus_one, dense_shape)->output(0);
+    labels = make_shared<v3::ScatterNDUpdate>(labels, decoded_indices, decoded_values);
 
-    // Compute label_lenght for each batch
-    auto minus_one_mask = make_shared<Equal>(decoded_values_dense, minus_one_value);
-    auto mask01 = make_shared<Select>(minus_one_mask,
-                                      make_shared<Constant>(element::i64, Shape{}, 1),
-                                      make_shared<Constant>(element::i64, Shape{}, 0));
-    auto label_length_axis = make_shared<Constant>(element::i64, Shape{}, 1);
-    auto label_length = make_shared<ReduceSum>(mask01, label_length_axis, false);
+    // compute label_lenght for each batch
+    auto minus_one_mask = make_shared<v1::Equal>(labels, minus_one);
+    auto mask01 = make_shared<v1::Select>(minus_one_mask,
+                                          make_shared<v0::Constant>(element::i64, Shape{}, 0),
+                                          make_shared<v0::Constant>(element::i64, Shape{}, 1));
+    auto reduce_axis = make_shared<v0::Constant>(element::i64, Shape{}, 1);
+    auto label_length = make_shared<v1::ReduceSum>(mask01, reduce_axis, false);
 
-    auto ctc_loss = make_shared<CTCLoss>(logits,
-                                         logit_length,
-                                         decoded_values_dense,
-                                         label_length,
-                                         preprocess_collapse_repeated,
-                                         ctc_merge_repeated);
+    auto ctc_loss = make_shared<v4::CTCLoss>(logits,
+                                             logit_length,
+                                             labels,
+                                             label_length,
+                                             preprocess_collapse_repeated,
+                                             ctc_merge_repeated);
     set_node_name(node.get_name(), ctc_loss);
     return {ctc_loss};
 }
diff --git a/tests/layer_tests/tensorflow_tests/test_tf_CTCLoss.py b/tests/layer_tests/tensorflow_tests/test_tf_CTCLoss.py
index 4f3cacb8369..5a9b024f509 100644
--- a/tests/layer_tests/tensorflow_tests/test_tf_CTCLoss.py
+++ b/tests/layer_tests/tensorflow_tests/test_tf_CTCLoss.py
@@ -1,12 +1,11 @@
 # Copyright (C) 2023 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+import numpy as np
 import pytest
-
+import tensorflow as tf
 from common.tf_layer_test_class import CommonTFLayerTest
 
-import numpy as np
-import tensorflow as tf
 
 # Testing operation CTCLoss
 # Documentation: https://www.tensorflow.org/api_docs/python/tf/raw_ops/CTCLoss
@@ -17,13 +16,7 @@ class TestCTCLoss(CommonTFLayerTest):
             inputs_dict[input] = np.random.randint(0, 5, inputs_dict[input]).astype(np.float32)
         return inputs_dict
 
-    def create_ctcloss_placeholder_const_net(self, inputs, targets, ir_version, use_new_frontend):
-        """
-            Tensorflow net                  IR net
-
-            Placeholder->CTCLoss    =>      Placeholder->Transpose->Convolution->Transpose #Need to replace by actual
-
-        """
+    def create_ctcloss_placeholder_const_net(self, inputs, targets):
         seq_lens = np.array([inputs[2]], dtype=np.int32)
         x = [targets]
 
@@ -40,7 +33,10 @@ class TestCTCLoss(CommonTFLayerTest):
         with tf.compat.v1.Session() as sess:
             tf_inputs = tf.compat.v1.placeholder(tf.float32, inputs, "inputs")
 
-            tf.raw_ops.CTCLoss(inputs = tf_inputs, labels_indices = indices, labels_values = vals, sequence_length = seq_lens)
+            ctc_loss = tf.raw_ops.CTCLoss(inputs=tf_inputs, labels_indices=indices, labels_values=vals,
+                                          sequence_length=seq_lens)
+            # compute exponent since CTCLoss value is -ln(prob)
+            tf.math.exp(-ctc_loss[0])
 
             tf.compat.v1.global_variables_initializer()
             tf_net = sess.graph_def
@@ -51,15 +47,15 @@ class TestCTCLoss(CommonTFLayerTest):
 
     # Reference values were copied from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py
     test_data = [
-        dict(inputs=[6,1,6], targets = [0, 1, 2, 1, 0]),
-        dict(inputs=[12,1,9], targets = [0, 1, 1, 0])
+        dict(inputs=[6, 1, 6], targets=[0, 1, 2, 1, 0]),
+        dict(inputs=[12, 1, 9], targets=[0, 1, 1, 0])
     ]
 
     @pytest.mark.parametrize("params", test_data)
+    @pytest.mark.precommit_tf_fe
     @pytest.mark.nightly
     def test_ctcloss_placeholder_const(self, params, ie_device, precision, ir_version, temp_dir,
-                                      use_new_frontend, use_old_api):
-        self._test(*self.create_ctcloss_placeholder_const_net(**params, ir_version=ir_version,
-                                                             use_new_frontend=use_new_frontend),
+                                       use_new_frontend, use_old_api):
+        self._test(*self.create_ctcloss_placeholder_const_net(**params),
                    ie_device, precision, ir_version, temp_dir=temp_dir,
-                   use_new_frontend=use_new_frontend, use_old_api=use_old_api)
+                   use_new_frontend=use_new_frontend, use_old_api=use_old_api, custom_eps=1e-3)