[PT FE] Add aten::atan2

src/frontends/pytorch/src/op/atan2.cpp

@@ -0,0 +1,100 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/add.hpp"
+#include "openvino/op/atan.hpp"
+#include "openvino/op/concat.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/convert.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/range.hpp"
+#include "openvino/op/select.hpp"
+#include "openvino/op/subtract.hpp"
+#include "pt_framework_node.hpp"
+#include "utils.hpp"
+
+using namespace ov::op;
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+OutputVector translate_atan2(const NodeContext& context) {
+    // Check whether inputs present
+    num_inputs_check(context, 2, 3);
+    // "aten::atan2.out(Tensor input,Tensor other, *,Tensor(a!) out) → Tensor(a!)"
+    Output<Node> y;
+    Output<Node> x;
+    // tie inputs together
+    std::tie(y, x) = get_inputs_with_promoted_types(context, 0, 1);
+    auto dummy_const = context.mark_node(ov::op::v0::Constant::create(element::f32, Shape({}), {0.5}))->output(0);
+    // align input types
+    align_eltwise_input_types(context, x, dummy_const, false, true);
+
+    // handle the first condition : x>0
+    auto div_y_x = context.mark_node(std::make_shared<v1::Divide>(y, x));
+    auto atan = context.mark_node(std::make_shared<v0::Atan>(div_y_x));
+    auto const_zero = v0::Constant::create(element::f32, Shape{}, {0});
+    auto result = atan->output(0);
+
+    // handle the second condition : x<0 && y>=0
+    auto const_pi = v0::Constant::create(element::f32, Shape{}, {std::atan(1.0) * 4});
+    // Same input type
+    x = context.mark_node(std::make_shared<v1::ConvertLike>(x, const_pi));
+    auto is_x_negative = context.mark_node(std::make_shared<v1::Less>(x, const_zero));
+    y = context.mark_node(std::make_shared<v1::ConvertLike>(y, const_zero));
+    auto y_non_negative = context.mark_node(std::make_shared<v1::GreaterEqual>(y, const_zero));
+    auto cond1 = context.mark_node(std::make_shared<v1::LogicalAnd>(is_x_negative, y_non_negative));
+    atan = context.mark_node(std::make_shared<v1::ConvertLike>(atan, const_pi));
+    auto atan_y_x_plus_pi = context.mark_node(std::make_shared<v1::Add>(atan, const_pi));
+    result = context.mark_node(std::make_shared<v1::Select>(cond1, atan_y_x_plus_pi, result));
+
+    // handle the third condition : x<0 && y<0
+    y = context.mark_node(std::make_shared<v1::ConvertLike>(y, const_zero));
+    auto is_y_negative = context.mark_node(std::make_shared<v1::Less>(y, const_zero));
+    auto cond2 = context.mark_node(std::make_shared<v1::LogicalAnd>(is_x_negative, is_y_negative));
+    y = context.mark_node(std::make_shared<v1::ConvertLike>(atan, const_pi));
+    auto atan_y_x_minus_pi = context.mark_node(std::make_shared<v1::Subtract>(atan, const_pi));
+    result = context.mark_node(std::make_shared<v1::Select>(cond2, atan_y_x_minus_pi, result));
+
+    // handle the fourth condition : x=0 && y>0
+    x = context.mark_node(std::make_shared<v1::ConvertLike>(x, const_zero));
+    auto is_x_zero = context.mark_node(std::make_shared<v1::Equal>(x, const_zero));
+    y = context.mark_node(std::make_shared<v1::ConvertLike>(y, const_zero));
+    auto is_y_positive = context.mark_node(std::make_shared<v1::Greater>(y, const_zero));
+    auto cond3 = context.mark_node(std::make_shared<v1::LogicalAnd>(is_x_zero, is_y_positive));
+    auto const_two = v0::Constant::create(element::f32, Shape{}, {2});
+    const_pi = context.mark_node(std::make_shared<v1::ConvertLike>(const_pi, const_two));
+    auto pi_div_two = context.mark_node(std::make_shared<v1::Divide>(const_pi, const_two));
+    result = context.mark_node(std::make_shared<v1::Select>(cond3, pi_div_two, result));
+
+    // handle the fifth condition : x=0 && y<0
+    auto cond4 = context.mark_node(std::make_shared<v1::LogicalAnd>(is_x_zero, is_y_negative));
+    auto const_minus_two = v0::Constant::create(element::f32, Shape{}, {-2});
+    const_minus_two = context.mark_node(std::make_shared<v1::ConvertLike>(const_minus_two, const_pi));
+    auto pi_div_minus_two = context.mark_node(std::make_shared<v1::Divide>(const_pi, const_minus_two));
+    result = context.mark_node(std::make_shared<v1::Select>(cond4, pi_div_two, result));
+
+    // check whether out tensor is given
+    if (!context.input_is_none(2) && context.get_input_size() == 3) {
+        context.mutate_input(2, result);
+    }
+
+    // when out tensor is not in input
+    return {result};
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -40,6 +40,7 @@ OP_CONVERTER(translate_argmax);
 OP_CONVERTER(translate_argmin);
 OP_CONVERTER(translate_as_strided);
 OP_CONVERTER(translate_as_tensor);
+OP_CONVERTER(translate_atan2);
 OP_CONVERTER(translate_avg_poolnd);
 OP_CONVERTER(translate_bool);
 OP_CONVERTER(translate_batch_norm);
@@ -359,6 +360,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::asinh_", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Asinh>>},
         {"aten::atan", op::optional_out<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Atan>, 1>},
         {"aten::atan_", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Atan>>},
+        {"aten::atan2", op::optional_out<op::translate_atan2, 2>},
         {"aten::atanh",
          op::optional_out<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Atanh>, 1>},
         {"aten::atanh_", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Atanh>>},

tests/layer_tests/pytorch_tests/test_atan2.py

@@ -0,0 +1,69 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0i
+
+import pytest
+import torch
+import math
+import numpy as np
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+class TestAtan2(PytorchLayerTest):
+    def _prepare_input(self, y, x, dtype1=None, dtype2=None):
+        inputs = [np.array(y).astype(dtype1), np.array(x).astype(dtype2)]
+        return inputs
+
+    def create_model(self, dtype1=None, dtype2=None use_out=False):
+        dtype_map = {
+            "float32": torch.float32,
+            "float64": torch.float64,
+            "int64": torch.int64,
+            "int32": torch.int32,
+            "int16": torch.int16,
+            "uint8": torch.uint8,
+            "int8": torch.int8,
+            }
+
+        class aten_atan2_out(torch.nn.Module):
+            def __init__(self, out) -> None:
+                super().__init__()
+                self.out = torch.empty(25, dtype=out)
+
+            def forward(self, y, x):
+                return torch.atan2(input = y, other = x, out=self.out)
+
+        class aten_atan2(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def forward(self, y, x):
+                return torch.atan2(input = y, other = x)
+
+        dtype1 = dtype_map.get(dtype1)
+        dtype2 = dtype_map.get(dtype2)
+
+        if use_out:
+            model_class = aten_atan2_out(dtype1)
+        else:
+            model_class = aten_atan2()
+
+
+        ref_net = None
+
+        return model_class, ref_net, "aten::atan2"
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("dtype1, dtype2", [(None, None), ("float32", "int32"), ("float64", "float64"), ("int32", "float64"), ("int64", "int16"), ("int8", "int8"), ("uint8", "uint8")])
+    @pytest.mark.parametrize(
+        "y, x", [(0, 1.5), (0, 0), (1.25, -5), (1, 10), (-1, -5.5), (-1, -5), (1.25, -5.5), (1.9, 2.9), [10, 9.9]]
+    )
+    @pytest.mark.parametrize("use_out", [False, True])
+    def test_atan2_with_out(self, dtype1, dtype2, use_out, y, x, ie_device, precision, ir_version):
+        self._test(
+            *self.create_model(dtype=dtype, use_out=use_out),
+            ie_device,
+            precision,
+            ir_version,
+            kwargs_to_prepare_input={"y": y, "x": x}
+        )