bug in MultiHeadAttention._compute_attention_mask
@ami-grynberg,
Thank you for the issue. Could you please provide the complete code to reproduce the issue. Also Could you please try to add the following lines at the beggining of the `call()` method in the **MultiHeadAttention** class:

```
if attention_mask is not None:
      attention_mask = tf.cast(tf.convert_to_tensor(attention_mask), tf.bool)
``` comment: 
    The bug does not involve a pre-defined attention_mask at the call level

Here is a short code to reproduce a good behavior and a buggy one:

import tensorflow.keras.layers as lyrs
import numpy as np

# --- good  rank == 2 ---#
  shp = [64,  16]
  vocab = np.prod(shp)
  x = np.arange(vocab).reshape(shp)

  emb = lyrs.Embedding(input_dim=vocab, output_dim=512, 
                       input_length=shp[-1],
                       mask_zero=True)
  mha = lyrs.MultiHeadAttention(num_heads=2, key_dim = 256)

  x = emb(x)
  x = mha(x,x,x)
  print(x)

# --- bug for rank > 2 ---#
  shp = [64, 12,  16]
  vocab = np.prod(shp)
  x = np.arange(vocab).reshape(shp)

  emb = lyrs.Embedding(input_dim=vocab, output_dim=512, 
                       input_length=shp[-1],
                       mask_zero=True)
  mha = lyrs.MultiHeadAttention(num_heads=2, key_dim = 256)

  x = emb(x) # no problem
  x = mha(x,x,x) # bug in computing attention_mask
  print(x) Please correct my buggy example:
shp = [64, 12, 16]
vocab = np.prod(shp)
x = np.arange(vocab).reshape(shp)

emb = lyrs.Embedding(input_dim=vocab, output_dim=512,
input_length=shp[-1],
mask_zero=True)
mha = lyrs.MultiHeadAttention(num_heads=2, key_dim = 256, attention_axes=(2,))

x = emb(x) # no problem
x = mha(x,x,x) # bug in computing attention_mask
print(x)
 @ami-grynberg,
Thank you for opening this issue. Development of keras moved to another [repository](https://github.com/keras-team/keras/issues). 

Could you please post this issue on keras-team/keras [repo](https://github.com/keras-team/keras/issues).
To know more please refer:
https://discuss.tensorflow.org/t/keras-project-moved-to-new-repository-in-https-github-com-keras-team-keras/1999
Thank you!
 This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60626">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60626">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

binary

### Tensorflow Version

tf.12.0

### Custom Code

No

### OS Platform and Distribution

windows 11

### Mobile device

_No response_

### Python version

3.9

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

There is a bug in MultiHeadAttention._compute_attention_mask when q, k or v ranks are more than 3.

Latest commit [67a8e6b](https://github.com/keras-team/keras/commit/67a8e6b96b1731ec732084f443345c70d87c31bf) on Jan 23

line 652:  auto_mask = query_mask[:, :, tf.newaxis]  # shape is [B, T, 1] 
fix:           auto_mask = tf.expand_dims(query_mask, axis =-1)

line 656:  mask = value_mask[:, tf.newaxis, :]  # shape is [B, 1, S]
fix:           mask = tf.expand_dims(key_mask, axis=-2)  # shape is [.., 1, S]

line 661: mask = key_mask[:, tf.newaxis, :]  # shape is [B, 1, S]
fix:          mask = tf.expand_dims(key_mask, axis=-2)  # shape is [.., 1, S]

### Standalone code to reproduce the issue

```shell
Submit any 4D input as query, key and value.
Error log below for input shapes [B, 12, 16, 128]
```


### Relevant log output

```shell
Message=Exception encountered when calling layer 'softmax' (type Softmax).

{{function_node __wrapped__AddV2_device_/job:localhost/replica:0/task:0/device:CPU:0}} Incompatible shapes: [64,12,1,16,16] vs. [64,12,1,12,16] [Op:AddV2]

Call arguments received by layer 'softmax' (type Softmax):
  • inputs=tf.Tensor(shape=(64, 12, 1, 16, 16), dtype=float32)
  • mask=tf.Tensor(shape=(64, 12, 1, 12, 16), dtype=bool)
```
</details>