[PT FE]: support aten::floor_divide_ and aten::movedim (#23446)
### Details: - *aten::floor_divide* - *aten::movedim* ### Tickets: - *ticket-id*
This commit is contained in:
parent
8ba1ae39ef
commit
f05391b6bc
|
|
@ -6,12 +6,17 @@
|
|||
|
||||
#include "openvino/frontend/pytorch/node_context.hpp"
|
||||
#include "openvino/op/add.hpp"
|
||||
#include "openvino/op/broadcast.hpp"
|
||||
#include "openvino/op/concat.hpp"
|
||||
#include "openvino/op/constant.hpp"
|
||||
#include "openvino/op/equal.hpp"
|
||||
#include "openvino/op/if.hpp"
|
||||
#include "openvino/op/non_zero.hpp"
|
||||
#include "openvino/op/not_equal.hpp"
|
||||
#include "openvino/op/range.hpp"
|
||||
#include "openvino/op/reshape.hpp"
|
||||
#include "openvino/op/scatter_elements_update.hpp"
|
||||
#include "openvino/op/shape_of.hpp"
|
||||
#include "openvino/op/unsqueeze.hpp"
|
||||
#include "utils.hpp"
|
||||
|
||||
|
|
@ -81,7 +86,47 @@ OutputVector translate_t(const NodeContext& context) {
|
|||
if_node->set_input(input, param_then, param_else);
|
||||
return {if_node->set_output(result_then, result_else)};
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
OutputVector translate_movedim(const NodeContext& context) {
|
||||
// aten::movedim.int(Tensor(a) self, int source, int destination) -> Tensor(a)
|
||||
// aten::movedim.intlist(Tensor(a) self, int[] source, int[] destination) -> Tensor(a)
|
||||
// based on https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/TensorShape.cpp#L3816
|
||||
num_inputs_check(context, 3, 3);
|
||||
auto x = context.get_input(0);
|
||||
auto src_dims = context.get_input(1);
|
||||
auto dst_dims = context.get_input(2);
|
||||
Output<Node> rank;
|
||||
std::tie(std::ignore, rank) = get_shape_rank(context, context.get_input(0), true);
|
||||
src_dims = normalize_axis(context, src_dims, rank);
|
||||
dst_dims = normalize_axis(context, dst_dims, rank);
|
||||
auto const_0 = context.mark_node(v0::Constant::create(element::i32, {}, {0}));
|
||||
auto const_1 = context.mark_node(v0::Constant::create(element::i32, {}, {1}));
|
||||
auto range = context.mark_node(std::make_shared<v4::Range>(const_0, rank, const_1, element::i32));
|
||||
auto dims_1d_shape = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1}));
|
||||
// operation accepts 0d and 1d source and destination, make them always 1d
|
||||
src_dims = context.mark_node(std::make_shared<v1::Reshape>(src_dims, dims_1d_shape, false));
|
||||
dst_dims = context.mark_node(std::make_shared<v1::Reshape>(dst_dims, dims_1d_shape, false));
|
||||
auto dims_shape = context.mark_node(std::make_shared<v3::ShapeOf>(src_dims, element::i32));
|
||||
auto minus_one_replaces = context.mark_node(std::make_shared<v1::Broadcast>(dims_1d_shape, dims_shape));
|
||||
// update position for the dim provided by user and mark used dims for source and destination as -1
|
||||
auto perm_dims = context.mark_node(std::make_shared<v3::ScatterElementsUpdate>(range, dst_dims, src_dims, const_0));
|
||||
auto src_perm_dims =
|
||||
context.mark_node(std::make_shared<v3::ScatterElementsUpdate>(range, src_dims, minus_one_replaces, const_0));
|
||||
auto dst_perm_dims =
|
||||
context.mark_node(std::make_shared<v3::ScatterElementsUpdate>(range, dst_dims, minus_one_replaces, const_0));
|
||||
// Remove the dims whose position we already know, the ones marked with -1 in previous step
|
||||
auto not_changed_src = context.mark_node(std::make_shared<v1::NotEqual>(src_perm_dims, dims_1d_shape));
|
||||
auto not_changed_dst = context.mark_node(std::make_shared<v1::NotEqual>(dst_perm_dims, dims_1d_shape));
|
||||
auto indices = context.mark_node(std::make_shared<v3::NonZero>(not_changed_dst, element::i32));
|
||||
auto updates = context.mark_node(std::make_shared<v3::NonZero>(not_changed_src, element::i32));
|
||||
// Update the position of the remaining dimensions. indices now contains the original position
|
||||
// updates contains the new position it will shifted to after considering the user inputs.
|
||||
indices = context.mark_node(std::make_shared<v1::Reshape>(indices, dims_1d_shape, false));
|
||||
updates = context.mark_node(std::make_shared<v1::Reshape>(updates, dims_1d_shape, false));
|
||||
auto scatter = std::make_shared<v3::ScatterElementsUpdate>(perm_dims, indices, updates, const_0);
|
||||
return {context.mark_node(std::make_shared<v1::Transpose>(x, scatter))};
|
||||
};
|
||||
|
||||
} // namespace op
|
||||
} // namespace pytorch
|
||||
|
|
|
|||
|
|
@ -135,6 +135,7 @@ OP_CONVERTER(translate_mean);
|
|||
OP_CONVERTER(translate_meshgrid);
|
||||
OP_CONVERTER(translate_min);
|
||||
OP_CONVERTER(translate_minimum);
|
||||
OP_CONVERTER(translate_movedim);
|
||||
OP_CONVERTER(translate_multinomial);
|
||||
OP_CONVERTER(translate_narrow);
|
||||
OP_CONVERTER(translate_native_multi_head_attention);
|
||||
|
|
@ -433,6 +434,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
|
|||
{"aten::floor", op::optional_out<op::translate_1to1_match_1_inputs<opset10::Floor>, 1>},
|
||||
{"aten::floor_", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Floor>>},
|
||||
{"aten::floor_divide", op::translate_floor_divide},
|
||||
{"aten::floor_divide_", op::inplace_op<op::translate_floor_divide>},
|
||||
{"aten::floordiv", op::translate_floor_divide},
|
||||
{"aten::fmod", op::translate_fmod},
|
||||
{"aten::frobenius_norm", op::translate_frobenius_norm},
|
||||
|
|
@ -520,6 +522,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
|
|||
{"aten::mish", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Mish>},
|
||||
{"aten::mish_", op::inplace_op<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Mish>>},
|
||||
{"aten::mm", op::translate_1to1_match_2_inputs<opset10::MatMul>},
|
||||
{"aten::movedim", op::translate_movedim},
|
||||
{"aten::mul", op::translate_mul},
|
||||
{"aten::mul_", op::translate_mul_},
|
||||
{"aten::multiply", op::translate_mul},
|
||||
|
|
|
|||
|
|
@ -41,6 +41,21 @@ class TestFloorDivide(PytorchLayerTest):
|
|||
|
||||
return aten_floor_divide(), ref_net, "aten::floor_divide"
|
||||
|
||||
def create_model_inplace(self):
|
||||
import torch
|
||||
|
||||
class aten_floor_divide_(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super(aten_floor_divide_, self).__init__()
|
||||
|
||||
def forward(self, input_tensor, other_tensor):
|
||||
return input_tensor.floor_divide_(other_tensor), input_tensor
|
||||
|
||||
ref_net = None
|
||||
|
||||
return aten_floor_divide_(), ref_net, "aten::floor_divide_"
|
||||
|
||||
|
||||
@pytest.mark.parametrize('input_tensor',
|
||||
([
|
||||
[5], [5, 5, 1], [1, 1, 5, 5],
|
||||
|
|
@ -65,6 +80,29 @@ class TestFloorDivide(PytorchLayerTest):
|
|||
self.other_tensor = other_tensor
|
||||
self._test(*self.create_model(), ie_device, precision, ir_version, trace_model=True, use_convert_model=True)
|
||||
|
||||
@pytest.mark.parametrize('input_tensor',
|
||||
([
|
||||
[5, 5, 5], [1, 1, 5, 5],
|
||||
]))
|
||||
@pytest.mark.parametrize('other_tensor',
|
||||
([
|
||||
np.array([0.5]).astype(np.float32), [5], [5, 1], [1, 5]
|
||||
]))
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
@pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
|
||||
reason='Ticket - 122715')
|
||||
def test_floor_divide_(self, input_tensor, other_tensor, ie_device, precision, ir_version):
|
||||
if type(input_tensor) is list:
|
||||
self.input_tensor = np.random.randn(*input_tensor).astype(np.float32)
|
||||
else:
|
||||
self.input_tensor = input_tensor
|
||||
if type(other_tensor) is list:
|
||||
self.other_tensor = np.random.randn(*other_tensor).astype(np.float32)
|
||||
else:
|
||||
self.other_tensor = other_tensor
|
||||
self._test(*self.create_model_inplace(), ie_device, precision, ir_version, trace_model=True, use_convert_model=True)
|
||||
|
||||
@pytest.mark.parametrize('input_data',
|
||||
[
|
||||
{ "tensor": [5], "low": 0, "high": 10 },
|
||||
|
|
|
|||
|
|
@ -50,6 +50,30 @@ class TestTranspose(PytorchLayerTest):
|
|||
self._test(*self.create_model(dim0, dim1, op_type), ie_device, precision, ir_version, trace_model=True)
|
||||
|
||||
|
||||
class TestMoveDim(PytorchLayerTest):
|
||||
def _prepare_input(self):
|
||||
return (np.random.randn(2, 3, 4, 5).astype(np.float32),)
|
||||
|
||||
def create_model(self, dim0, dim1):
|
||||
class aten_move_dim(torch.nn.Module):
|
||||
def __init__(self, dim0, dim1):
|
||||
super(aten_move_dim, self).__init__()
|
||||
self.dim0 = dim0
|
||||
self.dim1 = dim1
|
||||
|
||||
def forward(self, x):
|
||||
return torch.movedim(x, self.dim0, self.dim1)
|
||||
|
||||
ref_net = None
|
||||
|
||||
return aten_move_dim(dim0, dim1), ref_net, f"aten::movedim"
|
||||
|
||||
@pytest.mark.parametrize(("dim0", "dim1"), [[0, 1], [-1, 0], [2, -2], [3, 1], [3, 3], [[1, 2], [3, 0]], [[-4, 1], [1, -1]], [[1, 3, 2], [0, 1, 2 ]]])
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
def test_move_dim(self, dim0, dim1, ie_device, precision, ir_version):
|
||||
self._test(*self.create_model(dim0, dim1), ie_device, precision, ir_version, trace_model=True)
|
||||
|
||||
class TestTSmall(PytorchLayerTest):
|
||||
def _prepare_input(self, num_dims=2, input_dtype="float32"):
|
||||
shape = (2, 3)
|
||||
|
|
|
|||
Loading…
Reference in New Issue