fixed bug in testing strategy

ivy-llc · Sep 27, 2023 · c99bdbd · c99bdbd
1 parent 8f5ee33
commit c99bdbd
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/ivy_tests/test_ivy/test_functional/test_nn/test_layers.py b/ivy_tests/test_ivy/test_functional/test_nn/test_layers.py
@@ -225,6 +225,8 @@ def _mha_helper(draw, same_pre_embed_dim=False, batch_second=False):
         )
     )
 
+    if _self_attention and _qkv_same_dim:
+        _num_keys = _num_queries
     _static_shape = (_num_batches * num_heads, _num_keys, int(_embed_dim // num_heads))
     static_k = draw(
         st.one_of(
@@ -251,10 +253,7 @@ def _mha_helper(draw, same_pre_embed_dim=False, batch_second=False):
         )
     )
 
-    _mask_shape = (
-        _num_queries,
-        _num_queries if _self_attention and _qkv_same_dim else _num_keys,
-    )
+    _mask_shape = (_num_queries, _num_keys)
     if len(_batch_dim) and draw(st.booleans()):
         _mask_shape = (_num_batches * num_heads, *_mask_shape)
     attention_mask = draw(
@@ -272,7 +271,7 @@ def _mha_helper(draw, same_pre_embed_dim=False, batch_second=False):
         st.one_of(
             helpers.array_values(
                 dtype="bool",
-                shape=(*_batch_dim, _mask_shape[-1]),
+                shape=(*_batch_dim, _num_keys),
             ),
             st.none(),
         )