[PT FE]: support boolean input in aten::all (#22764)
### Details: - *item1* - *...* ### Tickets: - *ticket-id*
This commit is contained in:
parent
74ed7bee9d
commit
9bf497825e
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
Loading…
Reference in New Issue