[PT FE] fix invalid reshape shape after aten::index (#16821)

* [PT FE] fix invalid reshape shape after aten::index

* support aten::index_select

src/frontends/pytorch/src/op/index_select.cpp

@@ -0,0 +1,31 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/gather.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+OutputVector translate_index_select(const NodeContext& context) {
+    // aten::index_select(Tensor self, int dim, Tensor index) -> Tensor
+    // aten::index_select.out(Tensor self, int dim, Tensor index, *, Tensor(a!) out) -> Tensor(a!)
+    num_inputs_check(context, 3, 4);
+    auto x = context.get_input(0);
+    auto dim = context.get_input(1);
+    auto indicies = context.get_input(2);
+    auto gather = context.mark_node(std::make_shared<ov::op::v8::Gather>(x, indicies, dim));
+    if (!context.input_is_none(3)) {
+        context.mutate_input(3, gather);
+    }
+    return {gather};
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -59,6 +59,7 @@ OP_CONVERTER(translate_hardtanh);
 OP_CONVERTER(translate_if);
 OP_CONVERTER(translate_im2col);
 OP_CONVERTER(translate_index_put_);
+OP_CONVERTER(translate_index_select);
 OP_CONVERTER(translate_instance_norm);
 OP_CONVERTER(translate_int);
 OP_CONVERTER(translate_layer_norm);
@@ -230,6 +231,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
         {"aten::hardtanh_", op::inplace_op<op::translate_hardtanh>},
         {"aten::im2col", op::translate_im2col},
         {"aten::index_put_", op::inplace_op<op::translate_index_put_>},
+        {"aten::index_select", op::translate_index_select},
         {"aten::instance_norm", op::translate_instance_norm},
         {"aten::Int", op::translate_int},
         {"aten::IntImplicit", op::translate_int},

src/frontends/pytorch/src/transforms/aten_index_replacer.cpp

@@ -213,7 +213,7 @@ AtenIndexToSelect::AtenIndexToSelect() {
                 }
                 concat_dims.push_back(cum_adv_index_shape_tensor);
                 for (auto i : non_used_dims) {
-                    if (i < advanced_ids[i]) {
+                    if (i < advanced_ids[0]) {
                         continue;
                     }
                     concat_dims.push_back(input_dims->output(i));

tests/layer_tests/pytorch_tests/test_index.py

@@ -82,7 +82,7 @@ class TestIndexRange(PytorchLayerTest):
     def create_model(self):
         import torch
 
-        class aten_index_unsqueeze(torch.nn.Module):
+        class aten_index_arange(torch.nn.Module):
 
             def forward(self, x, y):
                 x = x.reshape(x.shape[0], -1)
@@ -90,7 +90,20 @@ class TestIndexRange(PytorchLayerTest):
 
         ref_net = None
 
-        return aten_index_unsqueeze(), ref_net, "aten::index"
+        return aten_index_arange(), ref_net, "aten::index"
+
+    def create_model2(self):
+        import torch
+
+        class aten_index_arange(torch.nn.Module):
+
+            def forward(self, x, y):
+                x = x.reshape(x.shape[0], x.shape[1], -1, 1)
+                return x[torch.arange(x.shape[0]), y]
+
+        ref_net = None
+
+        return aten_index_arange(), ref_net, "aten::index"
 
     @pytest.mark.nightly
     @pytest.mark.precommit
@@ -102,3 +115,13 @@ class TestIndexRange(PytorchLayerTest):
     def test_index_range(self, input_shape, idx, ie_device, precision, ir_version):
         self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={
                    "input_shape": input_shape, "idx": idx}, trace_model=True, dynamic_shapes=False)
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize(("input_shape", "idx"), (
+        ((1, 1), [0]),
+        ([2, 3], [1, 2]),
+        ([7, 8, 9], [1]),
+        ([2, 2, 3, 4], [0])))
+    def test_index_range_free_dims(self, input_shape, idx, ie_device, precision, ir_version):
+        self._test(*self.create_model2(), ie_device, precision, ir_version, kwargs_to_prepare_input={
+                   "input_shape": input_shape, "idx": idx}, trace_model=True, dynamic_shapes=False)

tests/layer_tests/pytorch_tests/test_index_select.py

@@ -0,0 +1,46 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+
+class TestIndexSelect(PytorchLayerTest):
+    def _prepare_input(self, index, out=False, dim=0):
+        import numpy as np
+        index = np.array(index).astype(np.int32)
+        input_data = np.random.randn(2, 3, 10, 10).astype(np.float32)
+        if not out:
+            return (input_data, index)
+        out = np.zeros_like(np.take(input_data, axis=dim, indices=index))
+        return (input_data, index, out)
+
+    def create_model(self, dim, out=False):
+        import torch
+
+        class aten_index_select(torch.nn.Module):
+            def __init__(self, dim, out=False):
+                super(aten_index_select, self).__init__()
+                self.dim = dim
+                if out:
+                    self.forward = self.forward_out
+
+            def forward(self, x, indices):
+                return torch.index_select(x, self.dim, indices)
+
+            def forward_out(self, x, indices, out):
+                return out, torch.index_select(x, self.dim, indices, out=out)
+
+        ref_net = None
+
+        return aten_index_select(dim, out), ref_net, "aten::index_select"
+
+    @pytest.mark.parametrize("dim", [0, 1, 2, 3, -1, -2, -3])
+    @pytest.mark.parametrize("indices", [[0, 1], [0], [1, 0]])
+    @pytest.mark.parametrize("out", [False, True])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_index_select(self, dim, out, indices, ie_device, precision, ir_version):
+        self._test(*self.create_model(dim, out), ie_device, precision, ir_version,
+                   kwargs_to_prepare_input={"index": indices, "out": out, "dim": dim})