jax

6617a0d1 - Expand `device_put` benchmarks to run with different numbers of arrays and input types

Expand `device_put` benchmarks to run with different numbers of arrays and input types For the upcoming batching changes for `device_put`, it is useful to benchmark `device_put` with varying numbers of arrays. PiperOrigin-RevId: 641716268