[PT FE]: support boolean input in aten::all (#22764)

### Details:
 - *item1*
 - *...*

### Tickets:
 - *ticket-id*
This commit is contained in:
Ekaterina Aidova 2024-02-12 18:10:44 +04:00 committed by GitHub
parent 74ed7bee9d
commit 9bf497825e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 73 additions and 18 deletions

View File

@ -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<ov::Node> 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<int64_t>(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<bool>(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<v0::Convert>(input_tensor, element::i32));
}
const auto all_nonzero = context.mark_node(std::make_shared<opset10::ReduceProd>(input_tensor, axes, keep_dims));
return {context.mark_node(std::make_shared<opset10::Convert>(all_nonzero, output_dtype))};
const auto all_nonzero = context.mark_node(std::make_shared<v1::ReduceProd>(input_tensor, axes, keep_dims));
auto result = context.mark_node(std::make_shared<v0::Convert>(all_nonzero, output_dtype));
if (!context.input_is_none(out_id)) {
context.mutate_input(out_id, result);
}
return {result};
};
} // namespace op

View File

@ -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})