[PT FE]: support aten::floor_divide_ and aten::movedim (#23446)

### Details:
 - *aten::floor_divide*
 - *aten::movedim*

### Tickets:
 - *ticket-id*
This commit is contained in:
Ekaterina Aidova 2024-03-14 13:29:38 +04:00 committed by GitHub
parent 8ba1ae39ef
commit f05391b6bc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 111 additions and 1 deletions

View File

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

View File

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

View File

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

View File

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