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