[PT FE] Support prim::DictConstruct on the output (#16894)

* Support dict on the output

* Preserve output order

src/frontends/pytorch/src/frontend.cpp

@@ -23,6 +23,7 @@
 #include "transforms/aten_index_put_replacer.hpp"
 #include "transforms/aten_index_replacer.hpp"
 #include "transforms/aten_stack_list_construct_replacer.hpp"
+#include "transforms/dict_resolver.hpp"
 #include "transforms/einsum_list_construct.hpp"
 #include "transforms/listconstruct_replacer.hpp"
 #include "transforms/min_max_prim_list_construct_replacer.hpp"
@@ -112,6 +113,7 @@ void FrontEnd::normalize(const std::shared_ptr<ov::Model>& model) const {
     manager.register_pass<ov::frontend::pytorch::pass::MinMaxPrimListConstructReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::StringEqualityReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::DecomposeListTupleResults>();
+    manager.register_pass<ov::frontend::pytorch::pass::DictResolver>();
     manager.register_pass<ov::pass::RemoveMultiSubGraphOpDanglingParamsResults>();
     manager.register_pass<ov::pass::ReverseShapeAndTypeInfer>();
 

src/frontends/pytorch/src/transforms/dict_resolver.cpp

@@ -0,0 +1,71 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "dict_resolver.hpp"
+
+#include "openvino/core/rt_info.hpp"
+#include "openvino/op/result.hpp"
+#include "openvino/op/util/framework_node.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace pass {
+
+using namespace ov::pass;
+using namespace ov::op;
+
+bool DictResolver::run_on_model(const std::shared_ptr<Model>& model) {
+    bool changed = false;
+    const auto results = model->get_results();
+    for (const auto& res : results) {
+        if (auto dict_construct_node = cast_fw_node(res->get_input_node_shared_ptr(0), "prim::DictConstruct")) {
+            const auto inputs = dict_construct_node->input_values();
+            if (inputs.size() % 2) {
+                // inputs must be divisible by 2
+                return false;
+            }
+            ResultVector new_outputs;
+            for (size_t i = 0; i < inputs.size(); i += 2) {
+                auto new_output = inputs.at(i + 1);
+                const auto name_node = inputs.at(i);
+                auto fw_node = std::dynamic_pointer_cast<ov::op::util::FrameworkNode>(name_node.get_node_shared_ptr());
+                if (!fw_node) {
+                    // odd inputs must contain constant strings encoded as fw nodes
+                    return false;
+                }
+                const auto attrs = fw_node->get_attrs();
+                if (attrs.find("string_value") == attrs.end()) {
+                    // fw node must contain string value
+                    return false;
+                }
+                auto name = attrs.at("string_value");
+                new_output.add_names({name});
+                new_outputs.push_back(std::make_shared<v0::Result>(new_output));
+            }
+            bool after_res = false;
+            for (const auto& result : results) {
+                if (after_res) {
+                    // To preserve output order remove results after this one and insert them after new outputs
+                    model->remove_result(result);
+                    new_outputs.push_back(result);
+                }
+                if (result == res) {
+                    after_res = true;
+                    model->remove_result(result);
+                }
+            }
+            model->add_results(new_outputs);
+            changed = true;
+        }
+    }
+    return changed;
+};
+
+}  // namespace pass
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/transforms/dict_resolver.hpp

@@ -0,0 +1,24 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+#include "openvino/pass/pass.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace pass {
+
+class DictResolver : public ov::pass::ModelPass {
+public:
+    OPENVINO_RTTI("ov::frontend::pytorch::pass::DictResolver");
+    bool run_on_model(const std::shared_ptr<Model>& model) override;
+};
+
+}  // namespace pass
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov