[CORE] Squeeze v15 reverse infer

src/common/transformations/src/transformations/common_optimizations/reverse_shape_and_type_infer.cpp

@@ -11,6 +11,7 @@
 #include "openvino/op/convert_like.hpp"
 #include "openvino/op/convolution.hpp"
 #include "openvino/op/deformable_convolution.hpp"
+#include "openvino/op/gather.hpp"
 #include "openvino/op/group_conv.hpp"
 #include "openvino/op/if.hpp"
 #include "openvino/op/parameter.hpp"
@@ -322,6 +323,21 @@ bool ov::pass::ReverseShapeAndTypeInfer::run_on_model(const std::shared_ptr<ov::
                 is_changed = true;
             }
             is_changed |= inherit_output_type(op, {0});
+        } else if (ov::as_type_ptr<ov::op::v0::Convert>(op)) {
+            is_changed |= inherit_output_shape(op, {0});
+            is_changed |= inherit_output_rank(op, {0});
+        } else if (auto gather_op = ov::as_type_ptr<ov::op::v8::Gather>(op)) {
+            is_changed |= inherit_output_type(op, {0});
+
+            const auto& output_shape = op->get_output_partial_shape(0);
+            const auto& data_shape = op->get_input_partial_shape(0);
+            const auto batch = gather_op->get_batch_dims();
+
+            if (op->get_input_size() > 1 && batch >= 0 && op->get_input_partial_shape(1).rank().is_dynamic()) {
+                op->get_input_tensor(1).m_partial_shape =
+                    ov::PartialShape::dynamic(output_shape.rank() - data_shape.rank() + batch + 1);
+                is_changed = true;
+            }
         }
     }
     return is_changed;

src/common/transformations/tests/common_optimizations/reverse_shape_and_type_infer.cpp

@@ -728,3 +728,125 @@ TEST_F(TransformationTestsF, TransposeWithConstantOrderReverseInfer2) {
         model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
     }
 }
+
+TEST_F(TransformationTestsF, GatherReverseInferIndicesRank) {
+    auto dyn = Dimension::dynamic();
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, 22, 333, 4444});
+        auto indices = std::make_shared<opset10::Parameter>(ov::element::i32, PartialShape::dynamic());
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {0});
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis);
+
+        // Concat is needed to produce static rank for indces
+        // Specify rank and type in one of Concat input to inherit in another
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 22, 333, 4444});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+        manager.register_pass<pass::ReverseShapeAndTypeInfer>();
+    }
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(element::f32, ov::PartialShape{1, 22, 333, 4444});
+        auto indices = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{dyn});
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {0});
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis);
+
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 22, 333, 4444});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+    }
+}
+
+TEST_F(TransformationTestsF, GatherReverseInferIndicesRankCustomBatchDims) {
+    auto dyn = Dimension::dynamic();
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{2, 5});
+        auto indices = std::make_shared<opset10::Parameter>(ov::element::i32, ov::PartialShape::dynamic());
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {1});
+        int64_t batch_dims = 1;
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis, batch_dims);
+
+        // Concat is needed to produce static rank for indces
+        // Specify rank and type in one of Concat input to inherit in another
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{2, 3});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+        manager.register_pass<pass::ReverseShapeAndTypeInfer>();
+    }
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{2, 5});
+        auto indices = std::make_shared<opset10::Parameter>(ov::element::i32, ov::PartialShape{dyn, dyn});
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {1});
+        int64_t batch_dims = 1;
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis, batch_dims);
+
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{2, 3});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+    }
+}
+
+TEST_F(TransformationTestsF, GatherReverseInferIndicesRankCustomBatchDims2) {
+    auto dyn = Dimension::dynamic();
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{2, 2, 5});
+        auto indices = std::make_shared<opset10::Parameter>(ov::element::i32, ov::PartialShape::dynamic());
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {2});
+        int64_t batch_dims = 2;
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis, batch_dims);
+
+        // Concat is needed to produce static rank for indces
+        // Specify rank and type in one of Concat input to inherit in another
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{2, 2, 3});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+        manager.register_pass<pass::ReverseShapeAndTypeInfer>();
+    }
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{2, 2, 5});
+        auto indices = std::make_shared<opset10::Parameter>(ov::element::i32, ov::PartialShape{dyn, dyn, dyn});
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {2});
+        int64_t batch_dims = 2;
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis, batch_dims);
+
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{2, 2, 3});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+    }
+}
+
+TEST_F(TransformationTestsF, GatherReverseInferIndicesRankCustomBatchDims3) {
+    auto dyn = Dimension::dynamic();
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{2, 1, 5, 4});
+        auto indices = std::make_shared<opset10::Parameter>(ov::element::i32, ov::PartialShape::dynamic());
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {2});
+        int64_t batch_dims = 1;
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis, batch_dims);
+
+        // Concat is needed to produce static rank for indces
+        // Specify rank and type in one of Concat input to inherit in another
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{2, 1, 3, 4});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+        manager.register_pass<pass::ReverseShapeAndTypeInfer>();
+    }
+    {
+        auto data = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{2, 1, 5, 4});
+        auto indices = std::make_shared<opset10::Parameter>(ov::element::i32, ov::PartialShape{dyn, dyn});
+        auto axis = ov::op::v0::Constant::create(element::i32, Shape{}, {2});
+        int64_t batch_dims = 1;
+        auto gather = std::make_shared<ov::op::v8::Gather>(data, indices, axis, batch_dims);
+
+        auto data2 = std::make_shared<opset10::Parameter>(element::f32, PartialShape{2, 1, 3, 4});
+        auto concat = std::make_shared<opset10::Concat>(OutputVector{gather, data2}, 1);
+        auto result = std::make_shared<opset10::Result>(concat);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data, indices, data2});
+    }
+}

src/common/transformations/tests/op_conversions/convert_squeeze15_downgrade_test.cpp

@@ -28,7 +28,7 @@ std::shared_ptr<ov::Model> create_v15_model(const IndicesMode indices_mode,
     ov::ParameterVector params = {data};
     std::shared_ptr<op::v15::Squeeze> squeeze;
     if (indices_mode == IndicesMode::NONE) {
-        squeeze = std::make_shared<ov::opset15::Squeeze>(data, allow_axis_skip);
+        squeeze = std::make_shared<ov::opset15::Squeeze>(data);
     } else if (indices_mode == IndicesMode::PARAM) {
         const auto& indices =
             std::make_shared<ov::opset15::Parameter>(ov::element::i32, PartialShape({data_shape.rank()}));

src/core/include/openvino/op/squeeze.hpp

@@ -45,6 +45,8 @@ class OPENVINO_API Squeeze : public util::SqueezeBase {
 public:
     OPENVINO_OP("Squeeze", "opset15");
 
+    static constexpr const char* output_rt_info_name = "rt_info_reverse_infer_output_rank";
+
     Squeeze();
     /// \brief Constructs a squeeze v15 operation.
     ///