Skip to content

Commit

Permalink
Remove assertion in MHA that last dims of Q and V need to be equal (#…
Browse files Browse the repository at this point in the history
…20340)

Via #20324.

This assertion was added in #19973, however re: #20324 it seems overly
restrictive. It was not covered by an existing test case.
  • Loading branch information
jeffcarp authored Oct 12, 2024
1 parent 7bd02bc commit 4b1e65e
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 15 deletions.
14 changes: 0 additions & 14 deletions keras/src/layers/attention/multi_head_attention.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,13 +211,6 @@ def build(
"""
key_shape = value_shape if key_shape is None else key_shape

if query_shape[-1] != value_shape[-1]:
raise ValueError(
"The last dimension of `query_shape` and `value_shape` "
f"must be equal, but are {query_shape[-1]}, {value_shape[-1]}. "
"Received: query_shape={query_shape}, value_shape={value_shape}"
)

if value_shape[1:-1] != key_shape[1:-1]:
raise ValueError(
"All dimensions of `value` and `key`, except the last one, "
Expand Down Expand Up @@ -604,13 +597,6 @@ def compute_output_shape(
if key_shape is None:
key_shape = value_shape

if query_shape[-1] != value_shape[-1]:
raise ValueError(
"The last dimension of `query_shape` and `value_shape` "
f"must be equal, but are {query_shape[-1]}, {value_shape[-1]}. "
"Received: query_shape={query_shape}, value_shape={value_shape}"
)

if value_shape[1:-1] != key_shape[1:-1]:
raise ValueError(
"All dimensions of `value` and `key`, except the last one, "
Expand Down
9 changes: 8 additions & 1 deletion keras/src/layers/attention/multi_head_attention_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,13 @@ def test_high_dim_attention(
(1, 1, 5, 2),
(3, 2),
),
(
"different_qv_last_dims",
(4, 2, 3, 8),
(4, 2, 3, 7),
(4, 2, 3, 8),
None,
),
)
def test_compute_output_shape(
self, query_dims, value_dims, key_dims, output_shape
Expand All @@ -130,7 +137,7 @@ def test_compute_output_shape(
self.assertEqual(output.shape, comp_output_shape)

@parameterized.named_parameters(
("query_value_dim_mismatch", (2, 4, 8), (2, 2, 7), 2),
("query_value_dim_mismatch", (2, 4, 8), (2, 2, 7), (2,)),
("key_value_dim_mismatch", (2, 4, 8), (2, 2, 8), (2, 1, 7)),
(
"key_value_dim_mismatch_high_dim",
Expand Down

0 comments on commit 4b1e65e

Please sign in to comment.