Change casting of float32 in top2 routing to allow for more stable bfloat16 training.
Change casting of float32 in top2 routing to allow for more stable bfloat16
training.

PiperOrigin-RevId: 342897841

tensorflow/mesh,tensorflow/mesh