Remove _multi_device_put, now that JAX has equivalent device_put_replicated. Also tweak module docstring.
Remove _multi_device_put, now that JAX has equivalent device_put_replicated.
Also tweak module docstring.

PiperOrigin-RevId: 354202824

google/trax,google/trax