def test_multi_head_attention_output_shape_as_int(self):
        """Test MultiHeadAttention with output_shape as an int."""
        mha = layers.MultiHeadAttention(num_heads=2, key_dim=16, output_shape=8)
        query = random.uniform((2, 4, 16))
        value = random.uniform((2, 4, 16))
        output = mha(query=query, value=value)

        assert output.shape == (
            2,
            4,
            8,
        ), f"Expected shape (2, 4, 8), got {output.shape}"
Test MultiHeadAttention with output_shape as an int.
keras/src/layers/attention/multi_head_attention_test.py