jax

b996073c - [JAX] Preserve `DeviceList` during a call to `batched_copy_array_to_devices_with_sharding()`

[JAX] Preserve `DeviceList` during a call to `batched_copy_array_to_devices_with_sharding()` This change makes `batched_copy_array_to_devices_with_sharding` take `Sequence[DeviceList]` for `devices` parameter, and avoids reconstructing these `DeviceList` objects. This improves the performance around `jax.device_put` when arrays are using many devices. In `jax/_src/array.py`, `_array_shard_arg()` was explicitly converting `Sequence[DeviceList]` into `Sequence[list[Device]]`. This conversion is now removed so that the code preserves `DeviceList`. Also, to plumb `DeviceList`, it uses `sharding._internal_device_list.addressable_device_list` instead of `sharding._addressable_device_assignment`. PiperOrigin-RevId: 774880618