Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put#5101

vfdev-5 wants to merge 1 commit intomaingoogle/flax:mainfrom

rewrite-prefetch_to_devicegoogle/flax:rewrite-prefetch_to_deviceCopy head branch name to clipboard