natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[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*
This commit is contained in:
Evgenya Nugmanova 2024-06-03 09:08:53 +04:00 committed by GitHub
parent 5a8d75d076
commit fc0b82c1be
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/core/shape_inference/include/scaled_dot_product_attention_shape_inference.hpp
View File

@ -73,9 +73,9 @@ std::vector<TRShape> 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 &&