Skip to content
Back to pull request #5101

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

Sign in to view logs

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

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

Test import standalone (3.12)

Test import standalone (3.12)
succeeded Dec 18, 2025 in 18s