Batch `pxla.shard_args` calls triggered by `jax.device_put`
Batch `pxla.shard_args` calls triggered by `jax.device_put`

With this change, one `jax.device_put` call now corresponds to one `device_put_p.bind()` instead of one per array. Immediately, this improves the performance of `jax.device_put(...)` with a large pytree by amortizing the calls to `pxla.shard_args`. Also, backends that implement efficient batch transfers (https://github.com/tensorflow/tensorflow/pull/69096) will batch device-to-device transfers across arrays in a pytree.

The api_benchmark indicates that `device_put` with a single array is slower by 10%, but `device_put` with 10 to 1000 arrays is ~30% faster. This seems acceptable since a single-array `device_put` is fast anyway (~45us on my workstation). If this becomes a problem, this can be addressed by rewriting part of `device_put` in C++.
