From 20bb450af21f03d57c7ae28eed24d3d88175dd4e Mon Sep 17 00:00:00 2001
From: Ivan Tikhonov <ivan.tikhonov@intel.com>
Date: Fri, 18 Aug 2023 10:23:59 +0400
Subject: [PATCH] Fix TransposeSinking when the data input is broadcasted
 (#19233)

---
 .../transpose_sinking/ts_utils.cpp            |  6 +--
 .../transpose_sinking/ts_common_test.cpp      | 40 +++++++++++++++++++
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/src/common/transformations/src/transformations/transpose_sinking/ts_utils.cpp b/src/common/transformations/src/transformations/transpose_sinking/ts_utils.cpp
index d28cf9ac256..522882e162e 100644
--- a/src/common/transformations/src/transformations/transpose_sinking/ts_utils.cpp
+++ b/src/common/transformations/src/transformations/transpose_sinking/ts_utils.cpp
@@ -257,9 +257,9 @@ NodeVector InsertOutputTransposes(const NodePtr& main_node, const TransposeInput
     NodeVector new_nodes;
 
     for (size_t i = 0; i < main_node->get_output_size(); ++i) {
-        auto new_transpose_const = std::make_shared<ov::op::v0::Constant>(transpose_element_type,
-                                                                          Shape{transpose_axis_order.size()},
-                                                                          transpose_axis_order);
+        auto aligned_order = AlignTransposeOrder(main_node->output(i), transpose_input_info);
+        auto new_transpose_const =
+            std::make_shared<ov::op::v0::Constant>(transpose_element_type, Shape{aligned_order.size()}, aligned_order);
         auto main_node_consumers = main_node->output(i).get_target_inputs();
         auto new_transpose = std::make_shared<ov::op::v1::Transpose>(main_node->output(i), new_transpose_const);
         for (auto& consumer : main_node_consumers) {
diff --git a/src/common/transformations/tests/transpose_sinking/ts_common_test.cpp b/src/common/transformations/tests/transpose_sinking/ts_common_test.cpp
index 336a89462c5..4bcfc42736b 100644
--- a/src/common/transformations/tests/transpose_sinking/ts_common_test.cpp
+++ b/src/common/transformations/tests/transpose_sinking/ts_common_test.cpp
@@ -432,6 +432,46 @@ auto test_forward_binary = []() {
 
 INSTANTIATE_TEST_SUITE_P(TransposeSinkingCommonBinaryForward, TSTestFixture, test_forward_binary());
 
+auto test_forward_binary_broadcasted = []() {
+    TestCase test_case;
+
+    // Initialize common attributes
+    test_case.transformation = CREATE_PASS_FACTORY(TSBinaryForward);
+    test_case.num_main_ops = {1, 10};
+    test_case.inputs_to_main = {
+        parameter(element::f32, {96, 55, 55}),
+        parameter(element::f32, {1, 1, 1, 1}),
+    };
+
+    // Test model description:
+    test_case.model.preprocess_inputs_to_main = {{set_transpose_for}, {{0}}};
+    test_case.model.main_op = binary_factories;
+    test_case.model.model_template = create_model;
+
+    // Reference model description:
+    auto new_transpose = [](const vector<size_t>& idxs, const OutputVector& out_vec) -> OutputVector {
+        OutputVector new_out_vec = out_vec;
+        shared_ptr<Node> order;
+        if (std::string(out_vec[idxs[0]].get_node_shared_ptr()->get_type_name()) == "PRelu") {
+            order = make_shared<Constant>(element::i32, Shape{3}, std::vector<int64_t>{2, 1, 0});
+        } else {
+            order = make_shared<Constant>(element::i32, Shape{4}, std::vector<int64_t>{0, 3, 2, 1});
+        }
+        new_out_vec[idxs[0]] = make_shared<Transpose>(out_vec[idxs[0]], order);
+        return new_out_vec;
+    };
+    test_case.model_ref.preprocess_inputs_to_main = {{new_transpose}, {{1}}};
+    test_case.model_ref.main_op = binary_factories;
+    test_case.model_ref.preprocess_outputs_of_main = {{new_transpose}, {{0}}};
+    test_case.model_ref.model_template = create_model;
+
+    return wrapper(test_case);
+};
+
+INSTANTIATE_TEST_SUITE_P(TransposeSinkingCommonBinaryForwardBroadcasted,
+                         TSTestFixture,
+                         test_forward_binary_broadcasted());
+
 auto test_forward_fq = []() {
     TestCase test_case;