jax
956f32e2 - Make `unreduced` transfers via device_put go via `batched_copy_array_to_devices_with_sharding` when there is a single process and no reshards.

Commit

13 days ago

Make `unreduced` transfers via device_put go via `batched_copy_array_to_devices_with_sharding` when there is a single process and no reshards. PiperOrigin-RevId: 940199458

References

#38861 - Make `unreduced` transfers via device_put go via `batched_copy_array_to_devices_with_sharding` when there is a single process and no reshards.

Author

yashk2810

Committer

Google-ML-Automation

Parents

ca9d73c4

jax 956f32e2 - Make `unreduced` transfers via device_put go via `batched_copy_array_to_devices_with_sharding` when there is a single process and no reshards.

jax
956f32e2 - Make `unreduced` transfers via device_put go via `batched_copy_array_to_devices_with_sharding` when there is a single process and no reshards.