From 9bf497825e5edb9bb8e960a062bb82f081db8945 Mon Sep 17 00:00:00 2001 From: Ekaterina Aidova Date: Mon, 12 Feb 2024 18:10:44 +0400 Subject: [PATCH] [PT FE]: support boolean input in aten::all (#22764) ### Details: - *item1* - *...* ### Tickets: - *ticket-id* --- src/frontends/pytorch/src/op/all.cpp | 40 ++++++++++++---- tests/layer_tests/pytorch_tests/test_all.py | 51 +++++++++++++++++---- 2 files changed, 73 insertions(+), 18 deletions(-) diff --git a/src/frontends/pytorch/src/op/all.cpp b/src/frontends/pytorch/src/op/all.cpp index 55c830df524..1c690df18a6 100644 --- a/src/frontends/pytorch/src/op/all.cpp +++ b/src/frontends/pytorch/src/op/all.cpp @@ -1,8 +1,10 @@ -// Copyright (C) 2018-2023 Intel Corporation +// Copyright (C) 2018-2024 Intel Corporation // SPDX-License-Identifier: Apache-2.0 // #include "openvino/frontend/pytorch/node_context.hpp" -#include "openvino/opsets/opset10.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/convert.hpp" +#include "openvino/op/reduce_prod.hpp" #include "utils.hpp" namespace ov { @@ -10,9 +12,21 @@ namespace frontend { namespace pytorch { namespace op { +using namespace ov::op; + OutputVector translate_all(const NodeContext& context) { - num_inputs_check(context, 1, 3); - const auto input_tensor = context.get_input(0); + // aten::all(Tensor self) -> Tensor + // aten::all.dim(Tensor self, int dim, bool keepdim=False) -> Tensor + // aten::all.all_out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + // aten::all.out(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + // aten::all.int(int[] self) -> bool + // aten::all.float(float[] self) -> bool + // aten::all.bool(bool[] self) -> bool + num_inputs_check(context, 1, 4); + auto input_tensor = context.get_input(0); + + auto num_inputs = context.get_input_size(); + size_t out_id; element::Type output_dtype = element::boolean; if (input_tensor.get_element_type() == element::u8) { @@ -21,18 +35,28 @@ OutputVector translate_all(const NodeContext& context) { bool keep_dims = false; ov::Output axes; - if (context.get_input_size() == 1) { + if (num_inputs < 3) { axes = get_axes_range(context, 0); + out_id = 1; } else { const auto dim = context.const_input(1); - axes = context.mark_node(opset10::Constant::create(element::i64, Shape{1}, {dim})); + axes = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {dim})); if (!context.input_is_none(2)) { keep_dims = context.const_input(2); } + out_id = 3; + } + // ReduceProd does not support bool input + if (input_tensor.get_element_type() == element::boolean) { + input_tensor = context.mark_node(std::make_shared(input_tensor, element::i32)); } - const auto all_nonzero = context.mark_node(std::make_shared(input_tensor, axes, keep_dims)); - return {context.mark_node(std::make_shared(all_nonzero, output_dtype))}; + const auto all_nonzero = context.mark_node(std::make_shared(input_tensor, axes, keep_dims)); + auto result = context.mark_node(std::make_shared(all_nonzero, output_dtype)); + if (!context.input_is_none(out_id)) { + context.mutate_input(out_id, result); + } + return {result}; }; } // namespace op diff --git a/tests/layer_tests/pytorch_tests/test_all.py b/tests/layer_tests/pytorch_tests/test_all.py index f4c2ed71ed4..d8b90c97d35 100644 --- a/tests/layer_tests/pytorch_tests/test_all.py +++ b/tests/layer_tests/pytorch_tests/test_all.py @@ -16,6 +16,13 @@ class aten_all_noparam(torch.nn.Module): def forward(self, input_tensor): return torch.all(input_tensor) +class aten_all_noparam_out(torch.nn.Module): + def __init__(self) -> None: + torch.nn.Module.__init__(self) + + def forward(self, input_tensor, out): + return torch.all(input_tensor, out=out), out + class aten_all(torch.nn.Module): def __init__(self, dim, keepdim) -> None: torch.nn.Module.__init__(self) @@ -32,9 +39,29 @@ class aten_all(torch.nn.Module): keepdim = self.keepdim ) +class aten_all_out(torch.nn.Module): + def __init__(self, dim, keepdim) -> None: + torch.nn.Module.__init__(self) + self.dim = dim + self.keepdim = keepdim + + def forward(self, input_tensor, out): + return torch.all( + input_tensor, + dim = self.dim, + out=out + ) if self.keepdim is None else torch.all( + input_tensor, + dim = self.dim, + keepdim = self.keepdim, + out=out + ), out + class TestAll(PytorchLayerTest): - def _prepare_input(self): - return (self.input_tensor,) + def _prepare_input(self, out=False): + if not out: + return (self.input_tensor,) + return (self.input_tensor, np.zeros_like(self.input_tensor, dtype=bool if self.input_tensor.dtype != np.uint8 else np.uint8)) @pytest.mark.parametrize("input_shape, d_type", [ (np.eye(5,5), np.int64), @@ -44,17 +71,19 @@ class TestAll(PytorchLayerTest): ([10, 13, 11], np.int64), ([8, 7, 6, 5, 4], np.int64), ([11, 11], np.uint8), - ([7, 7], np.uint8) + ([7, 7], np.uint8), + ([4, 4], bool) ]) + @pytest.mark.parametrize("out", [True, False]) @pytest.mark.nightly @pytest.mark.precommit - def test_all_noparams(self, input_shape, d_type, ie_device, precision, ir_version): + def test_all_noparams(self, input_shape, d_type, out, ie_device, precision, ir_version): if type(input_shape) is list: self.input_tensor = np.random.randint(0, 2, input_shape, dtype=d_type) else: self.input_tensor = input_shape - self._test(aten_all_noparam(), None, "aten::all", - ie_device, precision, ir_version, trace_model=True, freeze_model=False) + self._test(aten_all_noparam() if not out else aten_all_noparam_out(), None, "aten::all", + ie_device, precision, ir_version, trace_model=True, freeze_model=False, kwargs_to_prepare_input={"out": out}) @pytest.mark.parametrize("input_shape, d_type", [ (np.eye(5,5), np.int64), @@ -64,24 +93,26 @@ class TestAll(PytorchLayerTest): ([10, 13, 11], np.int64), ([8, 7, 6, 5, 4], np.int64), ([11, 11], np.uint8), - ([7, 7], np.uint8) + ([7, 7], np.uint8), + ([4, 4], bool) ]) @pytest.mark.parametrize("keepdim", [ True, False, None ]) + @pytest.mark.parametrize("out", [True, False]) @pytest.mark.nightly @pytest.mark.precommit @pytest.mark.xfail(condition=platform.system() in ('Darwin', 'Linux') and platform.machine() in ('arm', 'armv7l', 'aarch64', 'arm64', 'ARM64'), reason='Ticket - 122715') - def test_all(self, input_shape, d_type, keepdim, ie_device, precision, ir_version): + def test_all(self, input_shape, d_type, keepdim, out, ie_device, precision, ir_version): if type(input_shape) is list: self.input_tensor = np.random.randint(0, 2, input_shape, dtype=d_type) else: self.input_tensor = input_shape for dim in range(len(self.input_tensor.shape)): - self._test(aten_all(dim, keepdim), None, "aten::all", - ie_device, precision, ir_version, trace_model=True, freeze_model=False) + self._test(aten_all(dim, keepdim) if not out else aten_all_out(dim, keepdim), None, "aten::all", + ie_device, precision, ir_version, trace_model=True, freeze_model=False, kwargs_to_prepare_input={"out": out})