[PT FE]: Add aten::meshgrid + prim::ListUnpack transformation (#15195)

* Add meshgrid listunpack transformation

* Add case when indexing is not specified

* Fix typos

* Fix problem with 1 input execution & missing runtime_info

* Fix issue with meshgrid placed in loop body

* Add tests to precommit

* Apply sugestions from review

* Fix input 0

* Improve indexing attribute read

src/frontends/pytorch/src/op/meshgrid.cpp

@@ -0,0 +1,29 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "pt_framework_node.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+OutputVector translate_meshgrid(NodeContext& context) {
+    OutputVector inputs{context.get_input(0)};
+    std::string indexing = "ij";
+    if (!context.input_is_none(1)) {
+        indexing = context.const_input<std::string>(1);
+    }
+    auto node = std::make_shared<PtFrameworkNode>(context.get_decoder(), inputs);
+    auto attrs = node->get_attrs();
+    attrs["indexing"] = indexing;
+    node->set_attrs(attrs);
+    return {context.mark_node(node)};
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -62,6 +62,7 @@ OP_CONVERTER(translate_max);
 OP_CONVERTER(translate_masked_fill);
 OP_CONVERTER(translate_mean);
 OP_CONVERTER(translate_min);
+OP_CONVERTER(translate_meshgrid);
 OP_CONVERTER(translate_neg);
 OP_CONVERTER(translate_nonzero);
 OP_CONVERTER(translate_norm);
@@ -215,6 +216,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
         {"aten::max_pool3d", op::translate_max_poolnd},
         {"aten::max", op::translate_max},
         {"aten::mean", op::translate_mean},
+        {"aten::meshgrid", op::translate_meshgrid},
         {"aten::min", op::translate_min},
         {"aten::mm", op::translate_1to1_match_2_inputs<opset10::MatMul>},
         {"aten::bmm", op::translate_1to1_match_2_inputs<opset10::MatMul>},

src/frontends/pytorch/src/transforms/prim_list_unpack_replacer.cpp

@@ -137,6 +137,70 @@ PrimListUnpackReplacer::PrimListUnpackReplacer() {
 
             return true;
         }
+
+        if (auto meshgrid = cast_fw_node(input_node, "aten::meshgrid")) {
+            // Input - ListConstruct
+            auto meshgrid_input_node =
+                cast_fw_node(meshgrid->input_value(0).get_node_shared_ptr(), "prim::ListConstruct");
+            if (!meshgrid_input_node) {
+                return false;
+            }
+            NodeVector rt_copy_from{list_unpack, input_node, meshgrid_input_node};
+            OutputVector meshgrid_inputs;
+            for (auto& input : meshgrid_input_node->inputs()) {
+                meshgrid_inputs.push_back(input.get_source_output());
+            }
+
+            auto meshgrid_attrs = meshgrid->get_attrs();
+            if (meshgrid_attrs.find("indexing") == meshgrid_attrs.end()) {
+                // Check if "indexing" key is available in meshgrid attributes set in translation.
+                return false;
+            }
+            std::string indexing = meshgrid_attrs.at("indexing");
+            if (indexing != "ij" && indexing != "xy") {
+                // Check if indexing attribute has correct values.
+                return false;
+            }
+
+            if (indexing == "xy" && meshgrid_inputs.size() >= 2) {
+                std::swap(meshgrid_inputs[0], meshgrid_inputs[1]);
+            }
+            NodeVector cat_shapes{};
+            NodeVector reshapes{};
+            auto const_neg_1 = opset10::Constant::create(element::i64, Shape{1}, {-1});
+            auto const_1 = opset10::Constant::create(element::i64, Shape{1}, {1});
+            int input_idx = 0;
+            for (auto& input : meshgrid_inputs) {
+                auto reshaped_input = std::make_shared<opset10::Reshape>(input, const_neg_1, false);
+                auto shape = std::make_shared<opset10::ShapeOf>(reshaped_input);
+                cat_shapes.push_back(shape);
+                NodeVector cat_inputs{};
+                for (int i = 0; i < meshgrid_inputs.size(); i++) {
+                    cat_inputs.push_back(const_1);
+                }
+                cat_inputs[input_idx] = shape;
+                input_idx++;
+                auto input_cat = std::make_shared<opset10::Concat>(cat_inputs, 0);
+                auto reshape_cat = std::make_shared<opset10::Reshape>(reshaped_input, input_cat, false);
+                reshapes.push_back(reshape_cat);
+            }
+            auto cat = std::make_shared<opset10::Concat>(cat_shapes, 0);
+            NodeVector to_copy_rt{cat};
+            to_copy_rt.push_back(cat);
+            OutputVector outputs{};
+            for (auto& reshape : reshapes) {
+                auto out = std::make_shared<opset10::Broadcast>(reshape, cat, ov::op::BroadcastType::BIDIRECTIONAL);
+                to_copy_rt.push_back(out);
+                outputs.push_back(out);
+            }
+            if (indexing == "xy" && meshgrid_inputs.size() >= 2) {
+                std::swap(outputs[0], outputs[1]);
+            }
+            copy_runtime_info(rt_copy_from, to_copy_rt);
+            replace_node(list_unpack, outputs);
+            return true;
+        }
+
         if (auto shape_of = std::dynamic_pointer_cast<opset10::ShapeOf>(input_node)) {
             // case aten::size as input
             // Number of ListUnpack outputs should be equal to rank of input shape.

tests/layer_tests/pytorch_tests/test_listunpack.py

@@ -124,3 +124,81 @@ class TestListUnpack(PytorchLayerTest):
             precision,
             ir_version
         )
+
+class TestMeshgridListUnpack(PytorchLayerTest):
+    def _prepare_input(self):
+        return (
+            np.random.randn(3),
+            np.random.randn(5),
+            np.random.randn(7),
+            np.random.randn(11),
+        )
+
+    def create_model_meshgrid_listunpack_1_in(self, idx):
+        class prim_listunpack(torch.nn.Module):
+            def __init__(self, idx):
+                self.idx = idx
+                super(prim_listunpack, self).__init__()
+
+            def forward(self, in1, in2, in3, in4):
+                (
+                    a,
+                ) = torch.meshgrid(in1, indexing=self.idx)
+                return a
+
+        ref_net = None
+
+        return prim_listunpack(idx), ref_net, "prim::ListUnpack"
+
+    def create_model_meshgrid_listunpack_2_in(self, idx):
+        class prim_listunpack(torch.nn.Module):
+            def __init__(self, idx):
+                self.idx = idx
+                super(prim_listunpack, self).__init__()
+
+            def forward(self, in1, in2, in3, in4):
+                (
+                    a,
+                    b,
+                ) = torch.meshgrid(in1, in2, indexing=self.idx)
+                return a, b
+
+        ref_net = None
+
+        return prim_listunpack(idx), ref_net, "prim::ListUnpack"
+
+    def create_model_meshgrid_listunpack_3_in(self, idx):
+        class prim_listunpack(torch.nn.Module):
+            def __init__(self, idx):
+                self.idx = idx
+                super(prim_listunpack, self).__init__()
+
+            def forward(self, in1, in2, in3, in4):
+                a, b, c = torch.meshgrid(in1, in2, in3, indexing=self.idx)
+                return a, b, c
+
+        ref_net = None
+
+        return prim_listunpack(idx), ref_net, "prim::ListUnpack"
+
+    def create_model_meshgrid_listunpack_4_in(self, idx):
+        class prim_listunpack(torch.nn.Module):
+            def __init__(self, idx):
+                self.idx = idx
+                super(prim_listunpack, self).__init__()
+
+            def forward(self, in1, in2, in3, in4):
+                a, b, c, d = torch.meshgrid(in1, in2, in3, in4, indexing=self.idx)
+                return a, b, c, d
+
+        ref_net = None
+
+        return prim_listunpack(idx), ref_net, "prim::ListUnpack"
+
+    @pytest.mark.parametrize("idx", ["ij", "xy"])
+    @pytest.mark.parametrize("inp", [1, 2, 3, 4])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_meshgrid_listunpack(self, idx, inp, ie_device, precision, ir_version):
+        func = getattr(self, f"create_model_meshgrid_listunpack_{inp}_in")
+        self._test(*func(idx), ie_device, precision, ir_version)