[PT FE]: add more list construct transformations (#15614)

src/frontends/pytorch/src/transforms/listconstruct_replacer.cpp

@@ -9,8 +9,13 @@
 #include "openvino/op/broadcast.hpp"
 #include "openvino/op/concat.hpp"
 #include "openvino/op/constant.hpp"
+#include "openvino/op/equal.hpp"
 #include "openvino/op/reshape.hpp"
 #include "openvino/op/roll.hpp"
+#include "openvino/op/select.hpp"
+#include "openvino/op/shape_of.hpp"
+#include "openvino/op/tile.hpp"
+#include "openvino/op/transpose.hpp"
 #include "openvino/op/unsqueeze.hpp"
 #include "openvino/op/util/framework_node.hpp"
 #include "openvino/pass/pattern/matcher.hpp"
@@ -35,8 +40,25 @@ ListConstructReplacer::ListConstructReplacer() {
     auto roll_op = pattern::wrap_type<v7::Roll>({pattern::any_input(), list_construct, pattern::any_input()});
     auto broadcast_op = pattern::wrap_type<v3::Broadcast>({pattern::any_input(), list_construct});
     auto adapool_op = pattern::wrap_type<v8::AdaptiveAvgPool>({pattern::any_input(), list_construct});
-
-    auto lc_pattern = std::make_shared<pattern::op::Or>(OutputVector{reshape_op, roll_op, broadcast_op, adapool_op});
+    // replace list construct for aten::expand(tensor, prim::ListConstruct(shapes)) decomposition
+    //  shape_of + broadcast + equal + select
+    auto shape_of_op = pattern::wrap_type<v3::ShapeOf>({list_construct, pattern::any_input()});
+    auto equal_op = pattern::wrap_type<v1::Equal>({list_construct, pattern::any_input()});
+    auto select_op = pattern::wrap_type<v1::Select>({pattern::any_input(), pattern::any_input(), list_construct});
+    // replace list construct for aten::repeat(tensor,  prim::ListConstruct(shapes)))
+    // shape_of + broadcast + tile
+    auto tile_op = pattern::wrap_type<v0::Tile>({pattern::any_input(), list_construct});
+    // replace aten::permute(tensor, prim::ListConstruct)
+    auto transpose_op = pattern::wrap_type<v1::Transpose>({pattern::any_input(), list_construct});
+    auto lc_pattern = std::make_shared<pattern::op::Or>(OutputVector{reshape_op,
+                                                                     roll_op,
+                                                                     broadcast_op,
+                                                                     adapool_op,
+                                                                     shape_of_op,
+                                                                     equal_op,
+                                                                     select_op,
+                                                                     tile_op,
+                                                                     transpose_op});
 
     ov::matcher_pass_callback callback = [=](pattern::Matcher& m) {
         auto& pattern_map = m.get_pattern_value_map();

tests/layer_tests/pytorch_tests/test_expand.py

@@ -32,6 +32,30 @@ class TestExpand(PytorchLayerTest):
     def test_expand(self, dims, ie_device, precision, ir_version):
         self._test(*self.create_model(dims), ie_device, precision, ir_version)
 
+class TestExpandList(PytorchLayerTest):
+    def _prepare_input(self, broadcast_shape):
+        import numpy as np
+        return (np.random.randn(1, 3).astype(np.float32), np.random.randn(*broadcast_shape).astype(np.float32))
+
+    def create_model(self):
+        import torch
+
+        class aten_expand(torch.nn.Module):
+
+            def forward(self, x, y):
+                y_shape = y.shape
+                return x.expand([y_shape[0], y_shape[1]])
+
+        ref_net = None
+
+        return aten_expand(), ref_net, ["aten::expand", "prim::ListConstruct"]
+
+    @pytest.mark.parametrize("dims", [(3, 3), (2, 3), (1, 3), [4, 3]])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_expand(self, dims, ie_device, precision, ir_version):
+        self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={"broadcast_shape": dims})
+
 
 class TestExpandAs(PytorchLayerTest):
     def _prepare_input(self, input_shape, broadcast_shape):

tests/layer_tests/pytorch_tests/test_permute.py

@@ -29,5 +29,29 @@ class TestPermute(PytorchLayerTest):
     @pytest.mark.parametrize("order", [[0, 2, 3, 1], [0, 3, 1, 2]])
     @pytest.mark.nightly
     @pytest.mark.precommit
-    def test_relu(self, order, ie_device, precision, ir_version):
+    def test_permute(self, order, ie_device, precision, ir_version):
         self._test(*self.create_model(order), ie_device, precision, ir_version)
+
+class TestPermuteList(PytorchLayerTest):
+    def _prepare_input(self, permute_shape):
+        import numpy as np
+        return (np.random.randn(1, 3, 224, 224).astype(np.float32), np.random.randn(*permute_shape).astype(np.float32))
+
+    def create_model(self):
+        import torch
+
+        class aten_permute(torch.nn.Module):
+
+            def forward(self, x, y):
+                y_shape = y.shape
+                return torch.permute(x, [y_shape[0] - 1, y_shape[1] - 1, y_shape[2] - 1, y_shape[3] - 1])
+
+        ref_net = None
+
+        return aten_permute(), ref_net, ["aten::permute", "prim::ListConstruct"]
+
+    @pytest.mark.parametrize("order", [[1, 3, 4, 2], [1, 4, 2, 3]])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_permute(self, order, ie_device, precision, ir_version):
+        self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={"permute_shape": order})

tests/layer_tests/pytorch_tests/test_repeat.py

@@ -6,7 +6,7 @@ import pytest
 from pytorch_layer_test_class import PytorchLayerTest
 
 
-class TestRelu(PytorchLayerTest):
+class TestRepeat(PytorchLayerTest):
     def _prepare_input(self):
         import numpy as np
         return (np.random.randn(2, 10).astype(np.float32),)
@@ -29,5 +29,29 @@ class TestRelu(PytorchLayerTest):
     @pytest.mark.parametrize("repeats", [(4, 3), (1, 1), (1, 2, 3), (1, 2, 2, 3)])
     @pytest.mark.nightly
     @pytest.mark.precommit
-    def test_relu(self, repeats, ie_device, precision, ir_version):
+    def test_repeat(self, repeats, ie_device, precision, ir_version):
         self._test(*self.create_model(repeats), ie_device, precision, ir_version)
+
+class TestRepeatList(PytorchLayerTest):
+    def _prepare_input(self, repeats_shape):
+        import numpy as np
+        return (np.random.randn(2, 10).astype(np.float32), np.random.randn(*repeats_shape).astype(np.float32),)
+
+    def create_model(self):
+        import torch
+
+        class aten_repeat(torch.nn.Module):
+
+            def forward(self, x, y):
+                y_shape = y.shape
+                return x.repeat([y_shape[0], y_shape[1]])
+
+        ref_net = None
+
+        return aten_repeat(), ref_net, ["aten::repeat", "prim::ListConstruct"]
+
+    @pytest.mark.parametrize("repeats", [(4, 3), (1, 1), (1, 3, 3), (1, 2, 2, 3)])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_repeat(self, repeats, ie_device, precision, ir_version):
+        self._test(*self.create_model(), ie_device, precision, ir_version,  kwargs_to_prepare_input={"repeats_shape": repeats})