[PT FE] Implement override_all_inputs, override_all_outputs (#19642)

* Implement override_all_inputs, override_all_outputs

* Apply suggestions from code review

Co-authored-by: Piotr Krzemiński <piotrkrzeminski1234@gmail.com>

* Update src/frontends/pytorch/src/input_model.cpp

* Update src/frontends/pytorch/src/input_model.cpp

* Resolve problem with self input

* Update place.cpp

* Update place.cpp

* Fix build

---------

Co-authored-by: Piotr Krzemiński <piotrkrzeminski1234@gmail.com>

src/frontends/pytorch/include/openvino/frontend/pytorch/node_context.hpp

@@ -35,6 +35,9 @@ public:
           m_decoder_outputs(decoder->outputs()) {
         FRONT_END_GENERAL_CHECK(m_tensor_map != nullptr && m_external_parameters != nullptr &&
                                 m_mutated_tensors != nullptr && m_translate_session != nullptr);
+        for (size_t i = 0; i < m_decoder_inputs.size(); i++) {
+            m_inputs_is_none.push_back(decoder->input_is_none(i));
+        }
     }
 
     // Do not search for input in tensor map; try to access it as a constant of specified type T and return its value
@@ -48,7 +51,7 @@ public:
     // Search for input in tensor map and return an output port for already converted op
     // TODO: int due to base class uses it, but naturally it should be size_t for PT
     Output<Node> get_input(int index) const override {
-        FRONT_END_GENERAL_CHECK(!m_decoder->input_is_none(index), "Input is none with index: ", index);
+        FRONT_END_GENERAL_CHECK(!input_is_none(index), "Input is none with index: ", index);
         auto input = m_decoder_inputs.at(index);
         FRONT_END_GENERAL_CHECK(m_tensor_map->count(input), "No tensor corresponding input: ", input, " exist.");
         return m_tensor_map->at(input);
@@ -71,7 +74,7 @@ public:
     }
 
     bool input_is_none(size_t index) const {
-        return m_decoder->input_is_none(index);
+        return index >= m_inputs_is_none.size() || m_inputs_is_none.at(index);
     }
 
     Any get_output_type(size_t index) const {
@@ -147,6 +150,7 @@ private:
     TranslateSession* m_translate_session = nullptr;
     const std::vector<size_t> m_decoder_inputs;
     const std::vector<size_t> m_decoder_outputs;
+    std::vector<bool> m_inputs_is_none;
 };
 
 using CreatorFunction = std::function<ov::OutputVector(const ov::frontend::pytorch::NodeContext&)>;

src/frontends/pytorch/src/input_model.cpp

@@ -19,12 +19,7 @@ InputModel::InputModel(const std::shared_ptr<TorchDecoder>& model_decoder) : m_m
         for (const auto& name : in_place->get_names()) {
             m_name_to_place.emplace(name, std::dynamic_pointer_cast<frontend::Place>(in_place));
         }
-        auto type_any = simplified_type_interpret(m_model_decoder->get_input_type(i));
-        auto dtype = element::dynamic;
-        if (type_any.is<element::Type>()) {
-            dtype = type_any.as<element::Type>();
-        }
-        m_descriptors.emplace(inputs[i], PlaceDesc(dtype, m_model_decoder->get_input_shape(i)));
+        m_inputs.push_back(in_place);
     }
     const auto& outputs = m_model_decoder->outputs();
     for (size_t i = 0; i < outputs.size(); ++i) {
@@ -33,39 +28,24 @@ InputModel::InputModel(const std::shared_ptr<TorchDecoder>& model_decoder) : m_m
         for (const auto& name : out_place->get_names()) {
             m_name_to_place.emplace(name, std::dynamic_pointer_cast<frontend::Place>(out_place));
         }
-        auto type_any = simplified_type_interpret(m_model_decoder->get_output_type(i));
-        auto dtype = element::dynamic;
-        if (type_any.is<element::Type>()) {
-            dtype = type_any.as<element::Type>();
-        }
-        m_descriptors.emplace(outputs[i], PlaceDesc(dtype, m_model_decoder->get_output_shape(i)));
+        m_outputs.push_back(out_place);
     }
 }
 
 std::vector<ov::frontend::Place::Ptr> InputModel::get_inputs() const {
-    std::vector<ov::frontend::Place::Ptr> res;
-    for (const auto& input_idx : m_model_decoder->inputs()) {
-        auto place_it = m_name_to_place.find(std::to_string(input_idx));
-        FRONT_END_GENERAL_CHECK(place_it != m_name_to_place.end(), "Couldn't find Place for input.");
-        const auto& names = place_it->second->get_names();
-        if (input_idx == 0 && std::any_of(names.cbegin(), names.cend(), [](const std::string& n) {
+    if (m_inputs.size() > 0 && m_inputs[0]) {
+        // We need to remove "self" input to not confuse external users
+        const auto& names = m_inputs[0]->get_names();
+        if (std::any_of(names.cbegin(), names.cend(), [](const std::string& n) {
                 return n.find("self") != std::string::npos;
-            })) {
-            continue;
-        }
-        res.push_back(place_it->second);
+            }))
+            return std::vector<ov::frontend::Place::Ptr>(m_inputs.begin() + 1, m_inputs.end());
     }
-    return res;
+    return m_inputs;
 }
 
 std::vector<ov::frontend::Place::Ptr> InputModel::get_outputs() const {
-    std::vector<ov::frontend::Place::Ptr> res;
-    for (const auto& output_idx : m_model_decoder->outputs()) {
-        auto place_it = m_name_to_place.find(std::to_string(output_idx));
-        FRONT_END_GENERAL_CHECK(place_it != m_name_to_place.end(), "Couldn't find Place for output.");
-        res.push_back(place_it->second);
-    }
-    return res;
+    return m_outputs;
 }
 
 Place::Ptr InputModel::get_place_by_tensor_name(const std::string& tensor_name) const {
@@ -81,10 +61,7 @@ void InputModel::set_partial_shape(const Place::Ptr& place, const ov::PartialSha
                             "Provided place is invalid, only inputs are supported for setting shape.");
     auto pytorch_place = std::dynamic_pointer_cast<pytorch::Place>(place);
     FRONT_END_GENERAL_CHECK(pytorch_place, "Only place produced by PyTorch Frontend is supported");
-    auto it = m_descriptors.find(pytorch_place->get_tensor_index());
-    if (it != m_descriptors.end()) {
-        it->second.m_pshape = shape;
-    }
+    pytorch_place->m_pshape = shape;
 }
 
 ov::PartialShape InputModel::get_partial_shape(const Place::Ptr& place) const {
@@ -92,11 +69,7 @@ ov::PartialShape InputModel::get_partial_shape(const Place::Ptr& place) const {
     FRONT_END_GENERAL_CHECK(
         pytorch_place,
         "Provided place is invalid. Only place of input or output is supported by PyTorch Frontend.");
-    auto it = m_descriptors.find(pytorch_place->get_tensor_index());
-    if (it != m_descriptors.end()) {
-        return it->second.m_pshape;
-    }
-    return PartialShape::dynamic();
+    return pytorch_place->get_partial_shape();
 }
 
 void InputModel::set_element_type(const Place::Ptr& place, const ov::element::Type& type) {
@@ -104,10 +77,7 @@ void InputModel::set_element_type(const Place::Ptr& place, const ov::element::Ty
                             "Provided place is invalid, only inputs are supported for setting element type.");
     auto pytorch_place = std::dynamic_pointer_cast<pytorch::Place>(place);
     FRONT_END_GENERAL_CHECK(pytorch_place, "Only place produced by PyTorch Frontend is supported");
-    auto it = m_descriptors.find(pytorch_place->get_tensor_index());
-    if (it != m_descriptors.end()) {
-        it->second.m_type = type;
-    }
+    pytorch_place->m_type = type;
 }
 
 ov::element::Type InputModel::get_element_type(const Place::Ptr& place) const {
@@ -115,11 +85,7 @@ ov::element::Type InputModel::get_element_type(const Place::Ptr& place) const {
     FRONT_END_GENERAL_CHECK(
         pytorch_place,
         "Provided place is invalid. Only place of input or output is supported by PyTorch Frontend.");
-    auto it = m_descriptors.find(pytorch_place->get_tensor_index());
-    if (it != m_descriptors.end()) {
-        return it->second.m_type;
-    }
-    return element::dynamic;
+    return pytorch_place->get_element_type();
 }
 
 void InputModel::set_tensor_value(const Place::Ptr& place, const void* value) {
@@ -127,22 +93,84 @@ void InputModel::set_tensor_value(const Place::Ptr& place, const void* value) {
                             "Provided place is invalid, only inputs are supported for setting tensor value.");
     auto pytorch_place = std::dynamic_pointer_cast<pytorch::Place>(place);
     FRONT_END_GENERAL_CHECK(pytorch_place, "Only place produced by PyTorch Frontend is supported");
-    auto it = m_descriptors.find(pytorch_place->get_tensor_index());
+    const auto el_type = pytorch_place->m_type;
+    const auto p_shape = pytorch_place->m_pshape;
+    FRONT_END_GENERAL_CHECK(el_type.is_static() && p_shape.is_static(),
+                            "Shape and type must be statically defined before calling set_tensor_value");
+    auto const_node = ov::op::v0::Constant::create(el_type, p_shape.to_shape(), value);
+    const auto tensor_id = pytorch_place->get_tensor_index();
+    auto it = m_descriptors.find(tensor_id);
     if (it != m_descriptors.end()) {
-        auto el_type = it->second.m_type;
-        auto p_shape = it->second.m_pshape;
-        FRONT_END_GENERAL_CHECK(el_type.is_static() && p_shape.is_static(),
-                                "Shape and type must be statically defined before calling set_tensor_value");
-        it->second.m_value = ov::op::v0::Constant::create(el_type, p_shape.to_shape(), value);
+        it->second.m_value = const_node;
     } else {
-        FRONT_END_GENERAL_CHECK(false, "Place is not known.");
+        m_descriptors.emplace(tensor_id, PlaceDesc(const_node));
     }
+    // remove place from inputs
+    m_inputs.erase(std::remove_if(m_inputs.begin(),
+                                  m_inputs.end(),
+                                  [&](const Place::Ptr& p) {
+                                      return p->is_equal(place);
+                                  }),
+                   m_inputs.end());
+}
+
+void InputModel::override_all_outputs(const std::vector<Place::Ptr>& outputs) {
+    // Topology modification is not allowed, but order can be changed
+    auto all_outputs = std::all_of(outputs.cbegin(), outputs.cend(), [](const Place::Ptr& p) {
+        return p->is_output();
+    });
+    FRONT_END_GENERAL_CHECK(all_outputs, "Only initial outputs are supported by override_all_inputs.");
+    // Number of outputs can be lower then in original model
+    m_outputs = outputs;
+}
+
+void InputModel::override_all_inputs(const std::vector<Place::Ptr>& inputs) {
+    // Topology modification is not allowed, but order can be changed
+    auto all_inputs = std::all_of(inputs.cbegin(), inputs.cend(), [](const Place::Ptr& p) {
+        return p->is_input();
+    });
+    FRONT_END_GENERAL_CHECK(all_inputs, "Only initial inputs are supported by override_all_inputs.");
+    // We need to add back "self" input if it was in initial inputs
+    if (m_inputs.size() > 0 && m_inputs[0]) {
+        // We need to remove "self" input to not confuse external users
+        const auto& names = m_inputs[0]->get_names();
+        if (std::any_of(names.cbegin(), names.cend(), [](const std::string& n) {
+                return n.find("self") != std::string::npos;
+            })) {
+            FRONT_END_GENERAL_CHECK(inputs.size() == m_inputs.size() - 1,
+                                    "Number of inputs provided is incorrect. Graph modification is not supported for "
+                                    "this model. Expected number of inputs: ",
+                                    m_inputs.size() - 1,
+                                    " recieved ",
+                                    inputs.size());
+            auto self_place = m_inputs[0];
+            // Verify that no same place already in vector
+            auto no_self = std::all_of(inputs.cbegin(), inputs.cend(), [&](const Place::Ptr& p) {
+                return !p->is_equal(self_place);
+            });
+            FRONT_END_GENERAL_CHECK(no_self, "Unexpected input of 'self' was provided to override_all_inputs.");
+            m_inputs = std::vector<ov::frontend::Place::Ptr>{self_place};
+            m_inputs.insert(m_inputs.cend(), inputs.cbegin(), inputs.cend());
+            return;
+        }
+    }
+    FRONT_END_GENERAL_CHECK(inputs.size() == m_inputs.size(),
+                            "Number of inputs provided is incorrect. Graph modification is not supported for "
+                            "this model. Expected number of inputs: ",
+                            m_inputs.size(),
+                            " recieved ",
+                            inputs.size());
+    m_inputs = inputs;
 }
 
 const std::string& InputModel::decoder_type_name() const {
     return m_model_decoder->decoder_type_name();
 }
 
+std::shared_ptr<TorchDecoder> InputModel::get_decoder() const {
+    return m_model_decoder;
+}
+
 }  // namespace pytorch
 }  // namespace frontend
 }  // namespace ov

src/frontends/pytorch/src/input_model.hpp

@@ -17,18 +17,12 @@ class Place;
 class TorchDecoder;
 
 struct PlaceDesc {
-    PlaceDesc(const element::Type& type, const PartialShape& pshape)
-        : m_type(type),
-          m_pshape(pshape),
-          m_value(nullptr) {}
-    element::Type m_type;
-    PartialShape m_pshape;
+    PlaceDesc(std::shared_ptr<Node> value) : m_value(value) {}
     std::shared_ptr<Node> m_value;
 };
 
 class InputModel : public ov::frontend::InputModel {
     friend class ::ov::frontend::pytorch::TranslateSession;
-    friend class ::ov::frontend::pytorch::Place;
 
 public:
     explicit InputModel(const std::shared_ptr<TorchDecoder>& model_decoder);
@@ -41,11 +35,16 @@ public:
     void set_element_type(const frontend::Place::Ptr& place, const ov::element::Type& type) override;
     ov::element::Type get_element_type(const frontend::Place::Ptr& place) const override;
     void set_tensor_value(const frontend::Place::Ptr& place, const void* value) override;
+    void override_all_outputs(const std::vector<frontend::Place::Ptr>& outputs) override;
+    void override_all_inputs(const std::vector<frontend::Place::Ptr>& inputs) override;
     const std::string& decoder_type_name() const;
+    std::shared_ptr<TorchDecoder> get_decoder() const;
 
 private:
     std::shared_ptr<TorchDecoder> m_model_decoder;
     std::unordered_map<std::string, std::shared_ptr<frontend::Place>> m_name_to_place;
+    std::vector<std::shared_ptr<frontend::Place>> m_inputs;
+    std::vector<std::shared_ptr<frontend::Place>> m_outputs;
     std::unordered_map<size_t, PlaceDesc> m_descriptors;
 };
 

src/frontends/pytorch/src/node_context.cpp

@@ -54,7 +54,7 @@ std::shared_ptr<Node> NodeContext::mark_node(std::shared_ptr<Node> ov_node) cons
 }
 
 void NodeContext::mutate_input(size_t index, Output<Node> ov_output) const {
-    FRONT_END_GENERAL_CHECK(!m_decoder->input_is_none(index), "Input is none with index: ", index);
+    FRONT_END_GENERAL_CHECK(!input_is_none(index), "Input is none with index: ", index);
     auto input_id = m_decoder_inputs.at(index);
     FRONT_END_GENERAL_CHECK(m_tensor_map->count(input_id), "No tensor corresponding input: ", input_id, " exist.");
     m_translate_session->encode_tensor_name(ov_output, input_id, {m_decoder->get_input_debug_name(index)});

src/frontends/pytorch/src/place.cpp

@@ -8,6 +8,7 @@
 #include "openvino/frontend/exception.hpp"
 #include "openvino/frontend/pytorch/decoder.hpp"
 #include "openvino/util/log.hpp"
+#include "utils.hpp"
 
 namespace ov {
 namespace frontend {
@@ -15,25 +16,37 @@ namespace pytorch {
 
 Place::Place(const ov::frontend::InputModel& input_model, size_t tensor_index)
     : m_input_model(input_model),
-      m_tensor_index(tensor_index),
-      m_is_input(false),
-      m_is_output(false) {
+      m_tensor_index(tensor_index) {
     const auto im = dynamic_cast<const ov::frontend::pytorch::InputModel*>(&m_input_model);
     FRONT_END_GENERAL_CHECK(im, "PyTorch Place requires PyTorch InputModel class.");
-    const auto& inputs = im->m_model_decoder->inputs();
-    const auto& outputs = im->m_model_decoder->outputs();
+    auto decoder = im->get_decoder();
+    const auto& inputs = decoder->inputs();
+    const auto& outputs = decoder->outputs();
     auto in_it = std::find(inputs.begin(), inputs.end(), tensor_index);
     if (in_it != inputs.end()) {
         m_is_input = true;
-        const auto& signature_name =
-            im->m_model_decoder->get_input_signature_name(std::distance(inputs.begin(), in_it));
+        auto idx = std::distance(inputs.begin(), in_it);
+        const auto& signature_name = decoder->get_input_signature_name(idx);
         m_names.push_back(signature_name);
+
+        auto type_any = simplified_type_interpret(decoder->get_input_type(idx));
+        if (type_any.is<element::Type>()) {
+            m_type = type_any.as<element::Type>();
+        }
+        m_pshape = decoder->get_input_shape(idx);
     }
     auto out_it = std::find(outputs.begin(), outputs.end(), tensor_index);
     if (out_it != outputs.end()) {
         m_is_output = true;
-        const auto& debug_name = im->m_model_decoder->get_output_debug_name(std::distance(outputs.begin(), out_it));
+        auto idx = std::distance(outputs.begin(), out_it);
+        const auto& debug_name = decoder->get_output_debug_name(idx);
         m_names.push_back(debug_name);
+
+        auto type_any = simplified_type_interpret(decoder->get_output_type(idx));
+        if (type_any.is<element::Type>()) {
+            m_type = type_any.as<element::Type>();
+        }
+        m_pshape = decoder->get_output_shape(idx);
     }
     if (m_is_input && m_is_output) {
         OPENVINO_DEBUG << "[WARNING] Place " << tensor_index << " is input and output at a same time.";

src/frontends/pytorch/src/place.hpp

@@ -11,7 +11,12 @@ namespace ov {
 namespace frontend {
 namespace pytorch {
 
+class InputModel;
+
+// Place represents a tensor in the model
 class Place : public ov::frontend::Place {
+    friend class ::ov::frontend::pytorch::InputModel;
+
 public:
     Place(const ov::frontend::InputModel& input_model, size_t tensor_index);
 
@@ -32,7 +37,12 @@ public:
     size_t get_tensor_index() const {
         return m_tensor_index;
     }
-
+    element::Type get_element_type() const {
+        return m_type;
+    }
+    PartialShape get_partial_shape() const {
+        return m_pshape;
+    }
     bool is_equal_data(const Ptr& another) const override {
         const auto another_pt = dynamic_cast<ov::frontend::pytorch::Place*>(another.get());
         if (!another_pt) {
@@ -45,8 +55,10 @@ private:
     const ov::frontend::InputModel& m_input_model;
     const size_t m_tensor_index;
     std::vector<std::string> m_names;
-    bool m_is_input;
-    bool m_is_output;
+    element::Type m_type = element::dynamic;
+    PartialShape m_pshape;
+    bool m_is_input = false;
+    bool m_is_output = false;
 };
 
 }  // namespace pytorch

src/frontends/pytorch/src/translate_session.cpp

@@ -17,6 +17,7 @@
 #include "openvino/op/slice.hpp"
 #include "openvino/op/transpose.hpp"
 #include "openvino/util/log.hpp"
+#include "place.hpp"
 #include "pt_framework_node.hpp"
 #include "utils.hpp"
 
@@ -54,7 +55,7 @@ std::shared_ptr<ov::Model> TranslateSession::get_converted_model() {
 std::shared_ptr<ov::Model> TranslateSession::translate_graph(const ov::frontend::InputModel::Ptr& input_model) {
     auto pytorch_model = std::dynamic_pointer_cast<pytorch::InputModel>(input_model);
     FRONT_END_GENERAL_CHECK(pytorch_model != nullptr, "Invalid input model");
-    auto model = convert_pytorch_model(pytorch_model->m_model_decoder, {}, pytorch_model->m_descriptors);
+    auto model = convert_pytorch_model(pytorch_model->m_model_decoder, {}, pytorch_model);
     // First delete tensor indexes from outputs then resolve input names, otherwise Parameter->Result will fail
     for (auto& result : model->get_results()) {
         auto tensor_desc = result->input_value(0);
@@ -79,7 +80,7 @@ std::shared_ptr<ov::Model> TranslateSession::translate_graph(const ov::frontend:
 std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
     std::shared_ptr<TorchDecoder> pytorch_model,
     const TensorMap& external_tensor_map,
-    const std::unordered_map<size_t, PlaceDesc>& external_descriptors) {
+    const std::shared_ptr<pytorch::InputModel>& input_model) {
     // FIXME: don't use global variable to count inlined inputs, no we should use global counter because ID should be
     // unique for all new inlined inputs
     static size_t inlined_nodes_counter = 100000000;  // Suppose there are not graph with more than 10M nodes
@@ -90,36 +91,43 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
         auto tensor_map = std::make_shared<TensorMap>();  // tensor map of the current context
         auto mutated_tensors = std::make_shared<std::set<size_t>>();
 
-        //  Go over all pytorch_model inputs and register them in the tensor map:
-        auto inputs = pytorch_model->inputs();
-        for (size_t i = 0; i < inputs.size(); ++i) {
-            std::shared_ptr<Node> input_node;
-            element::Type type = element::dynamic;
-            PartialShape pshape;
-            auto desc = external_descriptors.find(inputs[i]);
-            if (desc != external_descriptors.end()) {
-                if (desc->second.m_value) {
-                    input_node = desc->second.m_value;
-                } else {
-                    pshape = desc->second.m_pshape;
-                    type = desc->second.m_type;
-                }
-            } else {
-                pshape = pytorch_model->get_input_shape(i);
+        if (input_model) {
+            // When we have input model we should use its inputs order to create Parameters
+            // We use m_inputs instead of get_inputs() because latter doesn't have "self" input
+            for (auto input_p : input_model->m_inputs) {
+                auto pytorch_place = std::dynamic_pointer_cast<pytorch::Place>(input_p);
+                FRONT_END_GENERAL_CHECK(pytorch_place, "Only place produced by PyTorch Frontend is supported.");
+                auto tensor_id = pytorch_place->get_tensor_index();
+                element::Type type = pytorch_place->get_element_type();
+                PartialShape pshape = pytorch_place->get_partial_shape();
+                auto parameter = std::make_shared<v0::Parameter>(type, pshape);
+                if (pytorch_place->get_names().size() > 0)
+                    parameter->set_friendly_name(pytorch_place->get_names().at(0));
+                encode_tensor_name(parameter->output(0), tensor_id);
+                parameters->push_back(parameter);
+                (*tensor_map)[tensor_id] = parameter;
+            }
+            // Add all tensors that were frozen
+            for (auto desc : input_model->m_descriptors) {
+                (*tensor_map)[desc.first] = desc.second.m_value;
+            }
+        } else {
+            // Go over all pytorch_model inputs and register them in the tensor map:
+            auto inputs = pytorch_model->inputs();
+            for (size_t i = 0; i < inputs.size(); ++i) {
+                element::Type type = element::dynamic;
+                PartialShape pshape = pytorch_model->get_input_shape(i);
                 auto type_any = simplified_type_interpret(pytorch_model->get_input_type(i));
                 // TODO: Use special API to set custom type specification
                 if (type_any.is<element::Type>()) {
                     type = type_any.as<element::Type>();
                 }
-            }
-            if (!input_node) {
                 auto parameter = std::make_shared<v0::Parameter>(type, pshape);
                 parameter->set_friendly_name(pytorch_model->get_input_signature_name(i));
                 encode_tensor_name(parameter->output(0), inputs.at(i), {pytorch_model->get_input_debug_name(i)});
                 parameters->push_back(parameter);
-                input_node = parameter;
+                (*tensor_map)[inputs.at(i)] = parameter;
             }
-            (*tensor_map)[inputs.at(i)] = input_node;
         }
 
         auto node_visitor = [&](std::shared_ptr<TorchDecoder> node) {
@@ -215,21 +223,36 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
         pytorch_model->visit_subgraph(node_visitor);
 
         ResultVector results;
-        for (size_t i = 0; i < pytorch_model->num_of_outputs(); ++i) {
-            size_t id = pytorch_model->output(i);
-            if (tensor_map->find(id) == tensor_map->end()) {
-                // Not found in this scope, adding Parameter to connect to external scope
-                auto parameter = std::make_shared<v0::Parameter>(element::dynamic, PartialShape::dynamic());
-                encode_tensor_name(parameter->output(0), id);
-                parameters->push_back(parameter);
-                (*tensor_map)[id] = parameter;
+        if (input_model) {
+            // For the case when we have InputModel we need to have same order as its outputs
+            for (auto output_p : input_model->get_outputs()) {
+                auto pytorch_place = std::dynamic_pointer_cast<pytorch::Place>(output_p);
+                FRONT_END_GENERAL_CHECK(pytorch_place, "Only place produced by PyTorch Frontend is supported.");
+                auto tensor_id = pytorch_place->get_tensor_index();
+                auto ov_output = tensor_map->at(tensor_id);
+                FRONT_END_GENERAL_CHECK(ov_output.get_names().size() > 0,
+                                        "Tensor doesn't have name, while it should have name: ",
+                                        tensor_id);
+                auto result = std::make_shared<v0::Result>(ov_output);
+                results.push_back(result);
+            }
+        } else {
+            for (size_t i = 0; i < pytorch_model->num_of_outputs(); ++i) {
+                size_t id = pytorch_model->output(i);
+                if (tensor_map->find(id) == tensor_map->end()) {
+                    // Not found in this scope, adding Parameter to connect to external scope
+                    auto parameter = std::make_shared<v0::Parameter>(element::dynamic, PartialShape::dynamic());
+                    encode_tensor_name(parameter->output(0), id);
+                    parameters->push_back(parameter);
+                    (*tensor_map)[id] = parameter;
+                }
+                auto ov_output = tensor_map->at(id);
+                FRONT_END_GENERAL_CHECK(ov_output.get_names().size() > 0,
+                                        "Tensor doesn't have name, while it should have name: ",
+                                        id);
+                auto result = std::make_shared<v0::Result>(ov_output);
+                results.push_back(result);
             }
-            auto ov_output = tensor_map->at(id);
-            FRONT_END_GENERAL_CHECK(ov_output.get_names().size() > 0,
-                                    "Tensor doesn't have name, while it should have name: ",
-                                    id);
-            auto result = std::make_shared<v0::Result>(ov_output);
-            results.push_back(result);
         }
 
         // Since parameters can be added we need to list all current parameters

src/frontends/pytorch/src/translate_session.hpp

@@ -30,10 +30,9 @@ public:
     /// context which is visible from nested model. Empty external_tensor_map is used as an indication that this is a
     /// main body conversion.
     /// \return fully converted OV Model
-    std::shared_ptr<Model> convert_pytorch_model(
-        std::shared_ptr<TorchDecoder> pytorch_model,
-        const TensorMap& external_tensor_map = {},
-        const std::unordered_map<size_t, PlaceDesc>& external_descriptors = {});
+    std::shared_ptr<Model> convert_pytorch_model(std::shared_ptr<TorchDecoder> pytorch_model,
+                                                 const TensorMap& external_tensor_map = {},
+                                                 const std::shared_ptr<pytorch::InputModel>& input_model = nullptr);
 
     /// \brief Returns backprop operations for direct operation
     Output<Node> get_backprop_op(const std::shared_ptr<TorchDecoder>& node,

tests/layer_tests/py_frontend_tests/test_torch_frontend.py

@@ -20,6 +20,11 @@ class aten_relu(torch.nn.Module):
         return x, torch.nn.functional.relu(x)
 
 
+class aten_multi_input_output(torch.nn.Module):
+    def forward(self, x, y, z):
+        return torch.nn.functional.relu(x), x * y, z / x
+
+
 def get_scripted_model(model):
     with torch.no_grad():
         model = torch.jit.script(model)
@@ -74,6 +79,40 @@ def test_pytorch_fe_set_input_value():
     assert len(om.get_parameters()) == 0
 
 
+def test_pytorch_fe_override_all_inputs():
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+
+    decoder = TorchScriptPythonDecoder(aten_multi_input_output())
+    fe_manager = FrontEndManager()
+    fe = fe_manager.load_by_framework("pytorch")
+    im = fe.load(decoder)
+    inputs = im.get_inputs()
+    assert len(inputs) == 3, "Model should have 3 inputs."
+    im.override_all_inputs([inputs[1], inputs[2], inputs[0]])
+    om = fe.convert(im)
+    assert om.get_parameters()[0].friendly_name == "y"
+    assert om.get_parameters()[1].friendly_name == "z"
+    assert om.get_parameters()[2].friendly_name == "x"
+
+
+def test_pytorch_fe_removes_input_set_value():
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+
+    decoder = TorchScriptPythonDecoder(aten_multi_input_output())
+    fe_manager = FrontEndManager()
+    fe = fe_manager.load_by_framework("pytorch")
+    im = fe.load(decoder)
+    assert len(im.get_inputs()) == 3, "Model should have 3 inputs."
+    place = im.get_place_by_tensor_name("y")
+    im.set_partial_shape(place, PartialShape([1]))
+    im.set_element_type(place, Type.f32)
+    im.set_tensor_value(place, np.random.randn(1).astype(np.float32))
+    assert len(im.get_inputs()) == 2, "Model should have 2 inputs."
+    om = fe.convert(im)
+    assert om.get_parameters()[0].friendly_name == "x"
+    assert om.get_parameters()[1].friendly_name == "z"
+
+
 def test_conversion_extension():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
@@ -108,7 +147,8 @@ def test_conversion_extension():
         inp = node.get_input(0)
         approximate = node.get_values_from_const_input(1)
         if approximate == "none":
-            f = ops.erf(ops.divide(inp, ops.constant(np.array([math.sqrt(2.0)], dtype=np.float32))))
+            f = ops.erf(ops.divide(inp, ops.constant(
+                np.array([math.sqrt(2.0)], dtype=np.float32))))
         elif approximate == "tanh":
             f = ops.tanh(ops.multiply(ops.constant(np.array([math.sqrt(2.0 / math.pi)], dtype=np.float32)),
                                       ops.add(inp, ops.multiply(ops.constant(np.array([0.044715], dtype=np.float32)),
@@ -128,7 +168,6 @@ def test_conversion_extension():
         div = ops.divide(exp, reduce_sum)
         return div.outputs()
 
-
     def convert_vector_norm(node: NodeContext):
         inp = node.get_input(0)
         ord = node.get_values_from_const_input(1)
@@ -142,14 +181,14 @@ def test_conversion_extension():
         rm = ops.reduce_max(ops.abs(inp), reduce_axes, False)
         return rm.outputs()
 
-
     fem = FrontEndManager()
     fe = fem.load_by_framework(framework="pytorch")
     assert fe
     fe.add_extension(ConversionExtension("aten::elu", convert_elu))
     fe.add_extension(ConversionExtension("aten::gelu", convert_gelu))
     fe.add_extension(ConversionExtension("aten::softmax", convert_softmax))
-    fe.add_extension(ConversionExtension("aten::linalg_vector_norm", convert_vector_norm))
+    fe.add_extension(ConversionExtension(
+        "aten::linalg_vector_norm", convert_vector_norm))
     input_model = fe.load(decoder)
     assert input_model
     converted_model = fe.convert(input_model)
@@ -177,7 +216,8 @@ def get_builtin_extensions_path():
         base_paths.append(repo_dir)
 
     for base_path in base_paths:
-        paths = glob.glob(os.path.join(base_path, "**", "*test_builtin_extensions*"), recursive=True)
+        paths = glob.glob(os.path.join(
+            base_path, "**", "*test_builtin_extensions*"), recursive=True)
         for path in paths:
             if re.search(r"(lib)?test_builtin_extensions.?\.(dll|so)", path):
                 return path
@@ -206,12 +246,14 @@ def test_op_extension():
     assert input_model
     converted_model = fe.convert(input_model)
     assert converted_model
-    assert [n.get_type_name() for n in converted_model.get_ordered_ops()] == ["Parameter", "Elu", "Result"]
+    assert [n.get_type_name() for n in converted_model.get_ordered_ops()] == [
+        "Parameter", "Elu", "Result"]
 
     fe.add_extension(get_builtin_extensions_path())
     converted_model = fe.convert(input_model)
     assert converted_model
-    assert [n.get_type_name() for n in converted_model.get_ordered_ops()] == ["Parameter", "CustomElu", "Result"]
+    assert [n.get_type_name() for n in converted_model.get_ordered_ops()] == [
+        "Parameter", "CustomElu", "Result"]
 
 
 def test_pytorch_telemetry():