[TF FE] Support Atan2 operation for TensorFlow models (openvinotoolki…

…t#21076) * Create atan2.cpp * Update common_op_table.hpp Adds atan2 operation to the list * Update op_table.cpp Adds atan2 to the list of BinaryOp. * Create test_tf_Atan2.py Adds test for Atan2. * Update src/frontends/tensorflow_common/src/op/atan2.cpp Co-authored-by: Roman Kazantsev <[email protected]> * Update src/frontends/tensorflow/src/op_table.cpp Co-authored-by: Roman Kazantsev <[email protected]> * Update atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp Co-authored-by: Roman Kazantsev <[email protected]> * Update atan2.cpp declares x and y , explicitly mentions conditions * Update atan2.cpp * Update atan2.cpp Adds versions. * Update atan2.cpp adds namespace * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update atan2.cpp Clang format * Update atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update src/frontends/tensorflow_common/src/op/atan2.cpp * Update test_tf_Atan2.py fixes the datatype of y * Update tests/layer_tests/tensorflow_tests/test_tf_Atan2.py * Update test_tf_Atan2.py shape to float * Update tests/layer_tests/tensorflow_tests/test_tf_Atan2.py * Update tests/layer_tests/tensorflow_tests/test_tf_Atan2.py * Update tests/layer_tests/tensorflow_tests/test_tf_Atan2.py * Update test_tf_Atan2.py Makes the range of input inclusive for x<0 and y<0. * Update src/frontends/tensorflow_common/src/op/atan2.cpp --------- Co-authored-by: Roman Kazantsev <[email protected]>
maxnick · Dec 6, 2023 · 55d7765 · 55d7765
1 parent 148daf0
commit 55d7765
Show file tree

Hide file tree

Showing 4 changed files with 121 additions and 1 deletion.
diff --git a/src/frontends/tensorflow/src/op_table.cpp b/src/frontends/tensorflow/src/op_table.cpp
@@ -102,6 +102,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
         // note: BinaryOp translator declaration for each op must to be added in binary_op.cpp file
         {"Add", CreatorFunction(translate_binary_op<opset8::Add>)},
         {"AddV2", CreatorFunction(translate_binary_op<opset8::Add>)},
+        {"Atan2", CreatorFunction(translate_atan2_op)},
         {"BitwiseAnd", CreatorFunction(translate_binary_op<opset13::BitwiseAnd>)},
         {"BitwiseOr", CreatorFunction(translate_binary_op<opset13::BitwiseOr>)},
         {"BitwiseXor", CreatorFunction(translate_binary_op<opset13::BitwiseXor>)},

diff --git a/src/frontends/tensorflow_common/include/common_op_table.hpp b/src/frontends/tensorflow_common/include/common_op_table.hpp
@@ -36,6 +36,7 @@ OP_CONVERTER(translate_add_n_op);
 OP_CONVERTER(translate_adjust_contrast_op);
 OP_CONVERTER(translate_arg_max_op);
 OP_CONVERTER(translate_arg_min_op);
+OP_CONVERTER(translate_atan2_op);
 OP_CONVERTER(translate_avg_pool_op);
 OP_CONVERTER(translate_batch_mat_mul_op);
 OP_CONVERTER(translate_batch_mat_mul_with_type_op);
@@ -173,4 +174,4 @@ OP_CONVERTER_NAMED(translate_unique_op);
 }  // namespace op
 }  // namespace tensorflow
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov
diff --git a/src/frontends/tensorflow_common/src/op/atan2.cpp b/src/frontends/tensorflow_common/src/op/atan2.cpp
@@ -0,0 +1,72 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include "common_op_table.hpp"
+#include "openvino/op/add.hpp"
+#include "openvino/op/atan.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/convert_like.hpp"
+#include "openvino/op/divide.hpp"
+#include "openvino/op/equal.hpp"
+#include "openvino/op/greater.hpp"
+#include "openvino/op/greater_eq.hpp"
+#include "openvino/op/less.hpp"
+#include "openvino/op/logical_and.hpp"
+#include "openvino/op/multiply.hpp"
+#include "openvino/op/select.hpp"
+#include "openvino/op/subtract.hpp"
+
+using namespace std;
+using namespace ov::op;
+
+namespace ov {
+namespace frontend {
+namespace tensorflow {
+namespace op {
+
+OutputVector translate_atan2_op(const NodeContext& node) {
+    default_op_checks(node, 2, {"Atan2"});
+    auto y = node.get_input(0);
+    auto x = node.get_input(1);
+
+    // handle the first condition : x>0
+    auto div_y_x = make_shared<v1::Divide>(y, x);
+    auto atan = make_shared<v0::Atan>(div_y_x);
+    auto const_zero = create_same_type_const_scalar<int32_t>(x, 0);
+    auto result = atan->output(0);
+
+    // handle the second condition : x<0 && y>=0
+    auto const_pi = create_same_type_const_scalar<double>(x, std::atan(1.0) * 4);
+    auto is_x_negative = make_shared<v1::Less>(x, const_zero);
+    auto y_non_negative = make_shared<v1::GreaterEqual>(y, const_zero);
+    auto cond1 = make_shared<v1::LogicalAnd>(is_x_negative, y_non_negative);
+    auto atan_y_x_plus_pi = make_shared<v1::Add>(atan, const_pi);
+    result = make_shared<v1::Select>(cond1, atan_y_x_plus_pi, result);
+
+    // handle the third condition : x<0 && y<0
+    auto is_y_negative = make_shared<v1::Less>(y, const_zero);
+    auto cond2 = make_shared<v1::LogicalAnd>(is_x_negative, is_y_negative);
+    auto atan_y_x_minus_pi = make_shared<v1::Subtract>(atan, const_pi);
+    result = make_shared<v1::Select>(cond2, atan_y_x_minus_pi, result);
+
+    // handle the fourth condition : x=0 && y>0
+    auto is_x_zero = make_shared<v1::Equal>(x, const_zero);
+    auto is_y_positive = make_shared<v1::Greater>(y, const_zero);
+    auto cond3 = make_shared<v1::LogicalAnd>(is_x_zero, is_y_positive);
+    auto const_two = create_same_type_const_scalar<int32_t>(x, 2);
+    auto pi_div_two = make_shared<v1::Divide>(const_pi, const_two);
+    result = make_shared<v1::Select>(cond3, pi_div_two, result);
+
+    // handle the fifth condition : x=0 && y<0
+    auto cond4 = make_shared<v1::LogicalAnd>(is_x_zero, is_y_negative);
+    auto const_minus_two = create_same_type_const_scalar<int32_t>(x, -2);
+    auto pi_div_minus_two = make_shared<v1::Divide>(const_pi, const_minus_two);
+    result = make_shared<v1::Select>(cond4, pi_div_two, result);
+
+    set_node_name(node.get_name(), result.get_node_shared_ptr());
+    return {result};
+}
+}  // namespace op
+}  // namespace tensorflow
+}  // namespace frontend
+}  // namespace ov
diff --git a/tests/layer_tests/tensorflow_tests/test_tf_Atan2.py b/tests/layer_tests/tensorflow_tests/test_tf_Atan2.py
@@ -0,0 +1,46 @@
+# Copyright (C) 2018-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
+
+
+class TestAtan2(CommonTFLayerTest):
+    def _prepare_input(self, inputs_info):
+        assert 'y' in inputs_info
+        assert 'x' in inputs_info
+        y_shape = inputs_info['y']
+        x_shape = inputs_info['x']
+        inputs_data = {}
+        inputs_data['y'] = np.random.rand(*y_shape).astype(self.input_type) - np.random.rand(*y_shape).astype(self.input_type)
+        inputs_data['x'] = np.random.rand(*x_shape).astype(self.input_type) - np.random.rand(*x_shape).astype(self.input_type)
+        return inputs_data
+
+    def create_atan2_net(self, input_shape, input_type):
+        self.input_type = input_type
+        tf.compat.v1.reset_default_graph()
+        # Create the graph and model
+        with tf.compat.v1.Session() as sess:
+            y = tf.compat.v1.placeholder(input_type, input_shape, 'y')
+            x = tf.compat.v1.placeholder(input_type, input_shape, 'x')
+            tf.raw_ops.Atan2(y=y, x=x)
+            tf.compat.v1.global_variables_initializer()
+            tf_net = sess.graph_def
+
+        return tf_net, None
+
+    test_data_basic = [
+        dict(input_shape=[1, 2], input_type=np.float32),
+        dict(input_shape=[2, 3, 4], input_type=np.float32),
+    ]
+
+    @pytest.mark.parametrize("params", test_data_basic)
+    @pytest.mark.precommit_tf_fe
+    @pytest.mark.nightly
+    def test_atan2_basic(self, params, ie_device, precision, ir_version, temp_dir,
+                               use_new_frontend, use_old_api):
+        self._test(*self.create_atan2_net(**params),
+                   ie_device, precision, ir_version, temp_dir=temp_dir,
+                   use_new_frontend=use_new_frontend, use_old_api=use_old_api)