Revises GMM Attention implementation.
Revises GMM Attention implementation.

This change includes:
- Changes shapes of some intermediate values.
- Fixes consistency on [target_batch] <-> [multiplier, source_batch] reshaping. This may remove a potential bug behavior.
- Removes duplicated tf.to_float() and pass dtype args directly to some fonctions.
- Adds some shape annotations.

PiperOrigin-RevId: 260431717

tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo