[PT FE] Support setting types/shapes for nested structures

src/bindings/python/src/pyopenvino/frontend/input_model.cpp

@@ -33,6 +33,18 @@ void regclass_frontend_InputModel(py::module m) {
                 :rtype: openvino.frontend.Place
              )");
 
+    im.def("get_place_by_input_index",
+           &ov::frontend::InputModel::get_place_by_input_index,
+           py::arg("input_idx"),
+           R"(
+                Returns a tensor place by an input index.
+
+                :param input_idx: Index of model input.
+                :type input_idx: int
+                :return: Tensor place corresponding to specified input index or nullptr.
+                :rtype: openvino.frontend.Place
+             )");
+
     im.def("get_place_by_operation_name",
            &ov::frontend::InputModel::get_place_by_operation_name,
            py::arg("operation_name"),

src/frontends/common/include/openvino/frontend/input_model.hpp

@@ -82,6 +82,11 @@ class FRONTEND_API InputModel {
     /// \return Tensor place corresponding to specified tensor name or nullptr if not exists
     virtual Place::Ptr get_place_by_tensor_name(const std::string& tensor_name) const;
 
+    /// \brief Returns a tensor place by an input index.
+    /// \param input_idx Index of model input
+    /// \return Tensor place corresponding to specified input index or nullptr
+    virtual Place::Ptr get_place_by_input_index(size_t input_idx) const;
+
     /// \brief Returns an operation place by an operation name following framework
     /// conventions, or nullptr if an operation with this name doesn't exist.
     /// \param operation_name Name of operation

src/frontends/common/src/input_model.cpp

@@ -33,6 +33,13 @@ Place::Ptr InputModel::get_place_by_tensor_name(const std::string& tensor_name)
     FRONTEND_RETURN_STATEMENT("get_place_by_tensor_name", m_actual->get_place_by_tensor_name(tensor_name))
 }
 
+Place::Ptr InputModel::get_place_by_input_index(size_t input_idx) const {
+    if (!m_actual) {
+        return {};
+    }
+    FRONTEND_RETURN_STATEMENT("get_place_by_input_index", m_actual->get_place_by_input_index(input_idx))
+}
+
 Place::Ptr InputModel::get_place_by_operation_name(const std::string& operation_name) const {
     if (!m_actual) {
         return {};

src/frontends/onnx/frontend/src/input_model.cpp

@@ -64,6 +64,14 @@ ov::frontend::Place::Ptr InputModel::get_place_by_tensor_name(const std::string&
     return nullptr;
 }
 
+ov::frontend::Place::Ptr InputModel::get_place_by_input_index(size_t input_idx) const {
+    const auto& inputs = get_inputs();
+    if (input_idx < inputs.size()) {
+        return inputs.at(input_idx);
+    }
+    return {};
+}
+
 ov::frontend::Place::Ptr InputModel::get_place_by_operation_name(const std::string& operation_name) const {
     if (m_editor->is_correct_and_unambiguous_node(operation_name)) {
         const auto node_index = m_editor->get_node_index(EditorNode{operation_name});

src/frontends/onnx/frontend/src/input_model.hpp

@@ -37,6 +37,7 @@ class InputModel : public ov::frontend::InputModel {
     std::vector<ov::frontend::Place::Ptr> get_inputs() const override;
     std::vector<ov::frontend::Place::Ptr> get_outputs() const override;
     ov::frontend::Place::Ptr get_place_by_tensor_name(const std::string& tensor_name) const override;
+    ov::frontend::Place::Ptr get_place_by_input_index(size_t input_idx) const override;
     ov::frontend::Place::Ptr get_place_by_operation_name(const std::string& operation_name) const override;
     ov::frontend::Place::Ptr get_place_by_operation_name_and_input_port(const std::string& operation_name,
                                                                         int input_port_index) override;

src/frontends/paddle/src/input_model.cpp

@@ -612,6 +612,14 @@ Place::Ptr InputModel::get_place_by_tensor_name(const std::string& tensorName) c
     return _impl->get_place_by_tensor_name(tensorName);
 }
 
+Place::Ptr InputModel::get_place_by_input_index(size_t input_idx) const {
+    const auto& inputs = get_inputs();
+    if (input_idx < inputs.size()) {
+        return inputs.at(input_idx);
+    }
+    return {};
+}
+
 void InputModel::override_all_outputs(const std::vector<Place::Ptr>& outputs) {
     _impl->override_all_outputs(outputs);
 }

src/frontends/paddle/src/input_model.hpp

@@ -25,6 +25,7 @@ class InputModel : public ov::frontend::InputModel {
     std::vector<Place::Ptr> get_inputs() const override;
     std::vector<Place::Ptr> get_outputs() const override;
     Place::Ptr get_place_by_tensor_name(const std::string& tensorName) const override;
+    Place::Ptr get_place_by_input_index(size_t input_idx) const override;
     void override_all_outputs(const std::vector<Place::Ptr>& outputs) override;
     void override_all_inputs(const std::vector<Place::Ptr>& inputs) override;
     void extract_subgraph(const std::vector<Place::Ptr>& inputs, const std::vector<Place::Ptr>& outputs) override;

src/frontends/pytorch/src/frontend.cpp

@@ -11,6 +11,7 @@
 #include "openvino/op/util/multi_subgraph_base.hpp"
 #include "openvino/pass/constant_folding.hpp"
 #include "openvino/util/log.hpp"
+#include "place.hpp"
 #include "pt_framework_node.hpp"
 #include "transformations/common_optimizations/push_constant_to_subgraph.hpp"
 #include "transformations/common_optimizations/remove_multi_subgraph_op_dangling_params.hpp"
@@ -113,12 +114,31 @@ std::string pack_detailed_failure_report(const std::map<std::string, std::string
         error_msg << '\n' << failed_ops_short.str();
     return error_msg.str();
 }
+
+void update_parameter_info(std::shared_ptr<ov::op::v0::Parameter>& param,
+                           const Place::Ptr& place,
+                           const std::shared_ptr<Model> ov_model) {
+    const auto& pt_place = std::dynamic_pointer_cast<pytorch::Place>(place);
+    FRONT_END_GENERAL_CHECK(pt_place, "Wrong place type.");
+    if (pt_place->get_element_type().is_static())
+        param->set_element_type(pt_place->get_element_type());
+    auto pshape = param->get_partial_shape();
+    FRONT_END_GENERAL_CHECK(PartialShape::merge_into(pshape, pt_place->get_partial_shape()),
+                            "Incompatible shape was requested for input ",
+                            ov_model->get_parameter_index(param),
+                            "\nOriginal shape: ",
+                            pshape,
+                            " update shape: ",
+                            pt_place->get_partial_shape());
+    param->set_partial_shape(pshape);
+}
 }  // namespace
 
 FrontEnd::FrontEnd() {}
 
 std::shared_ptr<Model> FrontEnd::convert(const ov::frontend::InputModel::Ptr& model) const {
-    FRONT_END_GENERAL_CHECK(std::dynamic_pointer_cast<pytorch::InputModel>(model), "Invalid input model");
+    auto pt_model = std::dynamic_pointer_cast<pytorch::InputModel>(model);
+    FRONT_END_GENERAL_CHECK(pt_model, "Invalid input model");
     std::map<std::string, CreatorFunction> supported_ops = get_supported_ops(model);
     std::shared_ptr<Model> converted_model;
     {
@@ -141,6 +161,38 @@ std::shared_ptr<Model> FrontEnd::convert(const ov::frontend::InputModel::Ptr& mo
     }
     bool is_conversion_successful = unconverted_ops.size() == 0 && norm_err.empty();
     FRONT_END_OP_CONVERSION_CHECK(is_conversion_successful, pack_detailed_failure_report(unconverted_ops, norm_err));
+
+    if (pt_model->m_requested_places.size() != 0) {
+        // Fake tensors mean that types were set to non-existent before conversion inputs.
+        // Here we resolve this. If input doesn't exist after conversion exception is raised.
+        const auto& inputs = pt_model->get_inputs();
+        auto parameters = converted_model->get_parameters();
+        for (size_t i = 0; i < inputs.size(); i++) {
+            FRONT_END_OP_CONVERSION_CHECK(i < parameters.size(),
+                                          "Type/shape was set to non-existent input. Converted model:\n",
+                                          converted_model);
+            update_parameter_info(parameters[i], inputs[i], converted_model);
+        }
+        for (const auto& fplace : pt_model->m_requested_places) {
+            const auto& pt_place = std::dynamic_pointer_cast<pytorch::Place>(fplace);
+            FRONT_END_GENERAL_CHECK(pt_place, "Wrong place type.");
+            if (fplace->get_names().size() == 0) {
+                // Fake place was set by input index
+                auto idx = pt_place->get_input_index();
+                FRONT_END_GENERAL_CHECK(idx < parameters.size(),
+                                        "Type/shape was set to non-existent input. Converted model:\n",
+                                        converted_model);
+                update_parameter_info(parameters[idx], fplace, converted_model);
+            } else {
+                auto input = converted_model->input(fplace->get_names().at(0));
+                auto param = std::dynamic_pointer_cast<ov::op::v0::Parameter>(input.get_node_shared_ptr());
+                FRONT_END_GENERAL_CHECK(param, "Input is not a Parameter.");
+                update_parameter_info(param, fplace, converted_model);
+            }
+        }
+        converted_model->validate_nodes_and_infer_types();
+    }
+
     return converted_model;
 }
 

src/frontends/pytorch/src/input_model.cpp

@@ -49,11 +49,27 @@ std::vector<ov::frontend::Place::Ptr> InputModel::get_outputs() const {
 }
 
 Place::Ptr InputModel::get_place_by_tensor_name(const std::string& tensor_name) const {
+    if (tensor_name.empty())
+        return {};
     auto place_it = m_name_to_place.find(tensor_name);
     if (place_it != m_name_to_place.end()) {
         return place_it->second;
+    } else {
+        // Return fake place that can be used to change shape or type of inputs that will exist after conversion
+        auto place = std::make_shared<pytorch::Place>(*this, tensor_name, 0);
+        return place;
+    }
+}
+
+Place::Ptr InputModel::get_place_by_input_index(size_t input_idx) const {
+    const auto& inputs = get_inputs();
+    if (input_idx < inputs.size()) {
+        return inputs.at(input_idx);
+    } else {
+        // Return fake place that can be used to change shape or type of inputs that will exist after conversion
+        auto place = std::make_shared<pytorch::Place>(*this, "", input_idx);
+        return place;
     }
-    return nullptr;
 }
 
 void InputModel::set_partial_shape(const Place::Ptr& place, const ov::PartialShape& shape) {
@@ -62,6 +78,12 @@ void InputModel::set_partial_shape(const Place::Ptr& place, const ov::PartialSha
     auto pytorch_place = std::dynamic_pointer_cast<pytorch::Place>(place);
     FRONT_END_GENERAL_CHECK(pytorch_place, "Only place produced by PyTorch Frontend is supported");
     pytorch_place->m_pshape = shape;
+    if (pytorch_place->m_is_fake &&
+        std::all_of(m_requested_places.cbegin(), m_requested_places.cend(), [&place](const Place::Ptr& p) {
+            return !p->is_equal(place);
+        })) {
+        m_requested_places.push_back(place);
+    }
 }
 
 ov::PartialShape InputModel::get_partial_shape(const Place::Ptr& place) const {
@@ -78,6 +100,12 @@ void InputModel::set_element_type(const Place::Ptr& place, const ov::element::Ty
     auto pytorch_place = std::dynamic_pointer_cast<pytorch::Place>(place);
     FRONT_END_GENERAL_CHECK(pytorch_place, "Only place produced by PyTorch Frontend is supported");
     pytorch_place->m_type = type;
+    if (pytorch_place->m_is_fake &&
+        std::all_of(m_requested_places.cbegin(), m_requested_places.cend(), [&place](const Place::Ptr& p) {
+            return !p->is_equal(place);
+        })) {
+        m_requested_places.push_back(place);
+    }
 }
 
 ov::element::Type InputModel::get_element_type(const Place::Ptr& place) const {

src/frontends/pytorch/src/input_model.hpp

@@ -12,6 +12,7 @@ namespace ov {
 namespace frontend {
 namespace pytorch {
 
+class FrontEnd;
 class TranslateSession;
 class Place;
 class TorchDecoder;
@@ -23,13 +24,15 @@ struct PlaceDesc {
 
 class InputModel : public ov::frontend::InputModel {
     friend class ::ov::frontend::pytorch::TranslateSession;
+    friend class ::ov::frontend::pytorch::FrontEnd;
 
 public:
     explicit InputModel(const std::shared_ptr<TorchDecoder>& model_decoder);
 
     std::vector<frontend::Place::Ptr> get_inputs() const override;
     std::vector<frontend::Place::Ptr> get_outputs() const override;
     frontend::Place::Ptr get_place_by_tensor_name(const std::string& tensor_name) const override;
+    frontend::Place::Ptr get_place_by_input_index(size_t input_idx) const override;
     void set_partial_shape(const frontend::Place::Ptr& place, const ov::PartialShape& shape) override;
     ov::PartialShape get_partial_shape(const frontend::Place::Ptr& place) const override;
     void set_element_type(const frontend::Place::Ptr& place, const ov::element::Type& type) override;
@@ -45,6 +48,7 @@ class InputModel : public ov::frontend::InputModel {
     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::vector<std::shared_ptr<frontend::Place>> m_requested_places;
     std::unordered_map<size_t, PlaceDesc> m_descriptors;
 };
 

src/frontends/pytorch/src/place.cpp

@@ -55,6 +55,18 @@ Place::Place(const ov::frontend::InputModel& input_model, size_t tensor_index)
     }
 }
 
+Place::Place(const ov::frontend::InputModel& input_model, const std::string& name, size_t input_index)
+    : m_input_model(input_model),
+      m_tensor_index(0),
+      m_is_fake(true),
+      m_input_index(input_index),
+      m_pshape(PartialShape::dynamic()),
+      m_type(element::dynamic),
+      m_is_input(true) {
+    if (!name.empty())
+        m_names = {name};
+}
+
 }  // namespace pytorch
 }  // namespace frontend
 }  // namespace ov

src/frontends/pytorch/src/place.hpp

@@ -19,6 +19,8 @@ class Place : public ov::frontend::Place {
 
 public:
     Place(const ov::frontend::InputModel& input_model, size_t tensor_index);
+    // Fake place of the input. If name is non-empty it is used as id of the input, otherwise input_index is used
+    Place(const ov::frontend::InputModel& input_model, const std::string& name, size_t input_index);
 
     ~Place() override = default;
 
@@ -50,13 +52,18 @@ class Place : public ov::frontend::Place {
         }
         return m_tensor_index == another_pt->get_tensor_index();
     }
+    size_t get_input_index() const {
+        return m_input_index;
+    }
 
 private:
     const ov::frontend::InputModel& m_input_model;
     const size_t m_tensor_index;
     std::vector<std::string> m_names;
-    element::Type m_type = element::dynamic;
+    bool m_is_fake = false;
+    const size_t m_input_index = 0;  // Only used for fake place
     PartialShape m_pshape;
+    element::Type m_type = element::dynamic;
     bool m_is_input = false;
     bool m_is_output = false;
 };

src/frontends/pytorch/src/translate_session.cpp

@@ -82,9 +82,10 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
         auto mutated_tensors = std::make_shared<std::set<size_t>>();
         std::vector<size_t> inserted_params;
 
-        if (input_model) {
+        if (input_model && input_model->m_requested_places.size() == 0) {
             // 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
+            // If there are fake places we don't need to use model inputs in that case
             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.");
@@ -98,10 +99,6 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
                 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();
@@ -120,6 +117,12 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
                 (*tensor_map)[inputs.at(i)] = parameter;
             }
         }
+        if (input_model) {
+            // Add all tensors that were frozen
+            for (auto& desc : input_model->m_descriptors) {
+                (*tensor_map)[desc.first] = desc.second.m_value;
+            }
+        }
 
         auto node_visitor = [&](std::shared_ptr<TorchDecoder> node) {
             // Explore all inputs of node. Node may refer to input value that hasn't been created in the current scope.

src/frontends/tensorflow/src/input_model.cpp

@@ -824,6 +824,14 @@ ov::frontend::Place::Ptr InputModel::get_place_by_tensor_name(const std::string&
     return _impl->get_place_by_tensor_name(tensorName);
 }
 
+ov::frontend::Place::Ptr InputModel::get_place_by_input_index(size_t input_idx) const {
+    const auto& inputs = get_inputs();
+    if (input_idx < inputs.size()) {
+        return inputs.at(input_idx);
+    }
+    return {};
+}
+
 void InputModel::override_all_outputs(const std::vector<ov::frontend::Place::Ptr>& outputs) {
     _impl->override_all_outputs(outputs);
 }

src/frontends/tensorflow/src/input_model.hpp

@@ -44,6 +44,7 @@ class InputModel : public ov::frontend::InputModel {
     std::vector<ov::frontend::Place::Ptr> get_inputs() const override;
     std::vector<ov::frontend::Place::Ptr> get_outputs() const override;
     ov::frontend::Place::Ptr get_place_by_tensor_name(const std::string& tensorName) const override;
+    ov::frontend::Place::Ptr get_place_by_input_index(size_t input_idx) const override;
     void override_all_outputs(const std::vector<ov::frontend::Place::Ptr>& outputs) override;
     void override_all_inputs(const std::vector<ov::frontend::Place::Ptr>& inputs) override;
     void extract_subgraph(const std::vector<ov::frontend::Place::Ptr>& inputs,

src/frontends/tensorflow_lite/src/input_model.cpp

@@ -406,6 +406,14 @@ ov::frontend::Place::Ptr InputModel::get_place_by_tensor_name(const std::string&
     return _impl->get_place_by_tensor_name(tensorName);
 }
 
+ov::frontend::Place::Ptr InputModel::get_place_by_input_index(size_t input_idx) const {
+    const auto& inputs = get_inputs();
+    if (input_idx < inputs.size()) {
+        return inputs.at(input_idx);
+    }
+    return {};
+}
+
 void InputModel::set_partial_shape(const Place::Ptr& place, const PartialShape& shape) {
     _impl->set_partial_shape(place, shape);
 }

src/frontends/tensorflow_lite/src/input_model.hpp

@@ -35,6 +35,7 @@ class InputModel : public ov::frontend::InputModel {
     std::vector<ov::frontend::Place::Ptr> get_inputs() const override;
     std::vector<ov::frontend::Place::Ptr> get_outputs() const override;
     ov::frontend::Place::Ptr get_place_by_tensor_name(const std::string& tensorName) const override;
+    ov::frontend::Place::Ptr get_place_by_input_index(size_t input_idx) const override;
 
     ///// Naming and annotation  /////
     void set_name_for_tensor(const Place::Ptr& tensor, const std::string& new_name) override;