diff --git a/src/frontends/pytorch/src/op/transpose.cpp b/src/frontends/pytorch/src/op/transpose.cpp index 3ea5ef17283..9b59bbebd16 100644 --- a/src/frontends/pytorch/src/op/transpose.cpp +++ b/src/frontends/pytorch/src/op/transpose.cpp @@ -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 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(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(src_dims, dims_1d_shape, false)); + dst_dims = context.mark_node(std::make_shared(dst_dims, dims_1d_shape, false)); + auto dims_shape = context.mark_node(std::make_shared(src_dims, element::i32)); + auto minus_one_replaces = context.mark_node(std::make_shared(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(range, dst_dims, src_dims, const_0)); + auto src_perm_dims = + context.mark_node(std::make_shared(range, src_dims, minus_one_replaces, const_0)); + auto dst_perm_dims = + context.mark_node(std::make_shared(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(src_perm_dims, dims_1d_shape)); + auto not_changed_dst = context.mark_node(std::make_shared(dst_perm_dims, dims_1d_shape)); + auto indices = context.mark_node(std::make_shared(not_changed_dst, element::i32)); + auto updates = context.mark_node(std::make_shared(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(indices, dims_1d_shape, false)); + updates = context.mark_node(std::make_shared(updates, dims_1d_shape, false)); + auto scatter = std::make_shared(perm_dims, indices, updates, const_0); + return {context.mark_node(std::make_shared(x, scatter))}; +}; } // namespace op } // namespace pytorch diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index 5bd891fa37a..82fdeb1c646 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -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 get_supported_ops_ts() { {"aten::floor", op::optional_out, 1>}, {"aten::floor_", op::inplace_op>}, {"aten::floor_divide", op::translate_floor_divide}, + {"aten::floor_divide_", op::inplace_op}, {"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 get_supported_ops_ts() { {"aten::mish", op::translate_1to1_match_1_inputs_with_fp32_type_alignment}, {"aten::mish_", op::inplace_op>}, {"aten::mm", op::translate_1to1_match_2_inputs}, + {"aten::movedim", op::translate_movedim}, {"aten::mul", op::translate_mul}, {"aten::mul_", op::translate_mul_}, {"aten::multiply", op::translate_mul}, diff --git a/tests/layer_tests/pytorch_tests/test_floor_divide.py b/tests/layer_tests/pytorch_tests/test_floor_divide.py index bd8827b2588..673cf53dfc0 100644 --- a/tests/layer_tests/pytorch_tests/test_floor_divide.py +++ b/tests/layer_tests/pytorch_tests/test_floor_divide.py @@ -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 }, diff --git a/tests/layer_tests/pytorch_tests/test_transpose.py b/tests/layer_tests/pytorch_tests/test_transpose.py index 5dec55ec59c..317a478ee79 100644 --- a/tests/layer_tests/pytorch_tests/test_transpose.py +++ b/tests/layer_tests/pytorch_tests/test_transpose.py @@ -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)