Fixes a shape bug in `MultiHeadedAttentionXL._AttenLogitsOneStep`
Fixes a shape bug in `MultiHeadedAttentionXL._AttenLogitsOneStep`

The tensor returned by `_AttenLogitsOneStep` is supposed to have shape `[S, B, N]`, but the einsum for the unsynced timestep and `skip_term_b` case produced a tensor with shape `BSN` instead.

PiperOrigin-RevId: 374252955

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