set `tf.expand_dims argument` `axis` to an integer instead of a list in the multi_head_attention model.
set `tf.expand_dims argument` `axis` to an integer instead of a list in the multi_head_attention model.

PiperOrigin-RevId: 380957913

keras-team/keras,keras-team/keras