Optimize `BatchedDevicePut` by batching GIL release
Optimize `BatchedDevicePut` by batching GIL release

The current `BatchedDevicePut` implementation interleaves Python work and C++ IFRT work, causing GIL to be repeatedly released/acquired for a single `jax.device_put()` call. This is particularly problematic for cases where there are a large number of devices.

This CL refactors `BatchedDevicePut` into two phases: `DevicePut` performs all Python work inline for formatting/casting the input object and returns a functor that performs C++ work. This allows us to change `BatchedDevicePut` to perform all Python work upfront, leaving all C++ work across all shards to be done in one GIL release.
