Use tpu_ops.all_to_all in mesh_tensorflow.
Use tpu_ops.all_to_all in mesh_tensorflow.

PiperOrigin-RevId: 211724278

tensorflow/tensor2tensor,vthorsteinsson/tensor2tensor,vthorsteinsson/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,vthorsteinsson/tensor2tensor,tensorflow/tensor2tensor,vthorsteinsson/tensor2tensor