[TRAX] Add a jax.device_get call to trax._save_replicated.
[TRAX] Add a jax.device_get call to trax._save_replicated.

By using device_get, trax can tell JAX that it intends to fetch all of the weights from the device, rather than synchronously transferring them one by one.

PiperOrigin-RevId: 251420028

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