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

Open

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

vfdev-5 wants to merge 1 commit into main from rewrite-prefetch_to_device

vfdev-5 marked this pull request as ready for review 44 days ago

cgarciae approved these changes on 2025-12-01

cgarciae added pull ready

cgarciae removed pull ready

cgarciae added pull ready

vfdev-5 force pushed from 5acf4e65 to 91e43b4e 15 days ago

Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replic…

65f71897

vfdev-5 force pushed from 91e43b4e to 65f71897 15 days ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone