[PT FE] Add aten::gcd (#21740)

* Add gcd and test

* constant with static

* exec_condition set to true

* boolean scalar, output condition, zero shape []

* removing unused headers

* more tests added

* code format

* Update src/frontends/pytorch/src/op/gcd.cpp

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>

* Update src/frontends/pytorch/src/op/gcd.cpp

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>

* Update src/frontends/pytorch/src/op/gcd.cpp

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>

* Update src/frontends/pytorch/src/op/gcd.cpp

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>

* ngraph header removed

* int64

* Add int64 support

* Update tests/layer_tests/pytorch_tests/test_gcd.py

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>

---------

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>
This commit is contained in:
Karan Jakhar 2024-01-19 00:25:25 +05:30 committed by GitHub
parent 6d4b2fe830
commit 103fb811e9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 194 additions and 0 deletions

View File

@ -0,0 +1,68 @@
// Copyright (C) 2018-2023 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "openvino/frontend/pytorch/node_context.hpp"
#include "openvino/op/constant.hpp"
#include "openvino/op/convert_like.hpp"
#include "openvino/op/loop.hpp"
#include "openvino/op/mod.hpp"
#include "openvino/op/not_equal.hpp"
#include "openvino/op/parameter.hpp"
#include "openvino/op/reduce_logical_or.hpp"
#include "openvino/op/select.hpp"
#include "openvino/openvino.hpp"
#include "utils.hpp"
namespace ov {
namespace frontend {
namespace pytorch {
namespace op {
using namespace ov::op;
OutputVector translate_gcd(const NodeContext& context) {
num_inputs_check(context, 2, 2);
auto x = context.get_input(0);
auto y = context.get_input(1);
align_eltwise_input_types(context, x, y, true);
auto zero_i32 = ov::op::v0::Constant::create(element::i32, Shape{}, {0});
auto trip_count = std::make_shared<v0::Constant>(element::i32, Shape{}, 1000);
auto exec_condition = std::make_shared<v0::Constant>(element::boolean, Shape{}, true);
auto loop = std::make_shared<op::v5::Loop>(trip_count, exec_condition);
auto x_input = std::make_shared<v0::Parameter>(x.get_element_type(), x.get_partial_shape());
auto y_input = std::make_shared<v0::Parameter>(y.get_element_type(), y.get_partial_shape());
x_input->set_element_type(x.get_element_type());
y_input->set_element_type(y.get_element_type());
auto zero = std::make_shared<ov::op::v1::ConvertLike>(zero_i32, x_input);
auto condition = std::make_shared<v1::NotEqual>(y_input, zero);
auto mod = std::make_shared<v1::Mod>(x_input, y_input);
auto new_x = std::make_shared<v1::Select>(condition, y_input, x_input);
auto new_y = std::make_shared<v1::Select>(condition, mod, zero);
auto reduced_condition = std::make_shared<v1::ReduceLogicalOr>(condition, zero);
auto body =
std::make_shared<Model>(OutputVector{new_x, new_y, reduced_condition}, ParameterVector{x_input, y_input});
loop->set_function(body);
loop->set_special_body_ports({-1, 2});
loop->set_merged_input(x_input, x, new_x);
loop->set_merged_input(y_input, y, new_y);
auto gcd_output = loop->get_iter_value(new_x, -1);
auto gcd_node = gcd_output.get_node_shared_ptr();
auto marked_gcd_node = context.mark_node(gcd_node);
return {marked_gcd_node};
}
} // namespace op
} // namespace pytorch
} // namespace frontend
} // namespace ov

View File

@ -84,6 +84,7 @@ OP_CONVERTER(translate_frobenius_norm);
OP_CONVERTER(translate_full);
OP_CONVERTER(translate_full_like);
OP_CONVERTER(translate_gather);
OP_CONVERTER(translate_gcd);
OP_CONVERTER(translate_gelu);
OP_CONVERTER(translate_get_attr);
OP_CONVERTER(translate_getitem);
@ -386,6 +387,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
{"aten::full", op::translate_full},
{"aten::full_like", op::translate_full_like},
{"aten::gather", op::translate_gather},
{"aten::gcd", op::translate_gcd},
{"aten::ge", op::translate_1to1_match_2_inputs_align_types<opset10::GreaterEqual>},
{"aten::gelu", op::translate_gelu},
{"aten::glu", op::translate_glu},

View File

@ -0,0 +1,124 @@
# Copyright (C) 2018-2023 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import numpy as np
import pytest
import torch
from pytorch_layer_test_class import PytorchLayerTest
class TestGcd(PytorchLayerTest):
def _prepare_input(self):
return self.input_data
def create_model_tensor_input(self):
class aten_gcd_tensor(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
def forward(self, tensor_a, tensor_b):
return torch.gcd(tensor_a, tensor_b)
ref_net = None
return aten_gcd_tensor(), ref_net, "aten::gcd"
def create_model_int_input(self):
class aten_gcd_int(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
def forward(self, int_a: int, int_b: int):
return torch.tensor(torch.gcd(int_a, int_b))
ref_net = None
return aten_gcd_int(), ref_net, "aten::gcd"
@pytest.mark.nightly
@pytest.mark.precommit
def test_gcd_int(self, ie_device, precision, ir_version):
self.input_data = (np.array(11, dtype=np.int32), np.array(17, dtype=np.int32))
self._test(
*self.create_model_int_input(),
ie_device,
precision,
ir_version,
use_convert_model=True,
trace_model=True
)
@pytest.mark.nightly
@pytest.mark.precommit
def test_gcd_tensor(self, ie_device, precision, ir_version):
self.input_data = (
np.array([14, 4, 12, 10, 3, 0], dtype=np.int32),
np.array([121, 2, 16, 0, 1, 8], dtype=np.int32),
)
self._test(
*self.create_model_tensor_input(),
ie_device,
precision,
ir_version,
use_convert_model=True,
trace_model=True
)
@pytest.mark.nightly
@pytest.mark.precommit
def test_gcd_int64(self, ie_device, precision, ir_version):
self.input_data = (np.array(11, dtype=np.int64), np.array(17, dtype=np.int64))
self._test(
*self.create_model_int_input(),
ie_device,
precision,
ir_version,
use_convert_model=True,
trace_model=True
)
@pytest.mark.nightly
@pytest.mark.precommit
def test_gcd_tensor64(self, ie_device, precision, ir_version):
self.input_data = (
np.array([14, 4, 12, 10, 3, 0], dtype=np.int64),
np.array([121, 2, 16, 0, 1, 8], dtype=np.int64),
)
self._test(
*self.create_model_tensor_input(),
ie_device,
precision,
ir_version,
use_convert_model=True,
trace_model=True
)
@pytest.mark.nightly
@pytest.mark.precommit
def test_gcd_int_diff_dtypes(self, ie_device, precision, ir_version):
self.input_data = (np.array(11, dtype=np.int64), np.array(17, dtype=np.int32))
self._test(
*self.create_model_int_input(),
ie_device,
precision,
ir_version,
use_convert_model=True,
trace_model=True
)
@pytest.mark.nightly
@pytest.mark.precommit
def test_gcd_tensor_diff_dtypes(self, ie_device, precision, ir_version):
self.input_data = (
np.array([14, 4, 12, 10, 3, 0], dtype=np.int64),
np.array([121, 2, 16, 0, 1, 8], dtype=np.int32),
)
self._test(
*self.create_model_tensor_input(),
ie_device,
precision,
ir_version,
use_convert_model=True,
trace_model=True
)