From fc0b82c1be18f793b41e894981a061b90c31a1e1 Mon Sep 17 00:00:00 2001 From: Evgenya Nugmanova Date: Mon, 3 Jun 2024 09:08:53 +0400 Subject: [PATCH] [SDPA] relax restriction on L from query and attention_mask being equal (#24745) ### Details: - *relax restriction on L from query and attention_mask being equal* ### Tickets: - *CVS-129000* --- .../scaled_dot_product_attention_shape_inference.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/core/shape_inference/include/scaled_dot_product_attention_shape_inference.hpp b/src/core/shape_inference/include/scaled_dot_product_attention_shape_inference.hpp index 7b38bc8c10d..b6b60cfad3a 100644 --- a/src/core/shape_inference/include/scaled_dot_product_attention_shape_inference.hpp +++ b/src/core/shape_inference/include/scaled_dot_product_attention_shape_inference.hpp @@ -73,9 +73,9 @@ std::vector shape_infer(const ScaledDotProductAttention* op, const auto& attention_mask_rank = attention_mask.rank(); if (attention_mask_rank.is_static() && attention_mask_rank != 0) { const auto& attention_mask_rank_len = attention_mask_rank.get_length(); - bool attention_mask_input_correctness = attention_mask_rank_len >= 2 && - DimType::merge(l_dim, l_dim, *(attention_mask.end() - 2)) && - DimType::merge(s_dim, s_dim, *(attention_mask.end() - 1)); + bool attention_mask_input_correctness = + attention_mask_rank_len >= 2 && DimType::broadcast_merge(l_dim, l_dim, *(attention_mask.end() - 2)) && + DimType::broadcast_merge(s_dim, s_dim, *(attention_mask.end() - 1)); if (attention_mask_rank_len >= 3) { attention_mask_input_correctness = attention_mask_input_correctness &&