Skip to content

Added tests for flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate#5101

Open
vfdev-5 wants to merge 1 commit intomainfrom
rewrite-prefetch_to_device
Open

Added tests for flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate#5101
vfdev-5 wants to merge 1 commit intomainfrom
rewrite-prefetch_to_device

Conversation

@vfdev-5
Copy link
Copy Markdown
Collaborator

@vfdev-5 vfdev-5 commented Nov 19, 2025

Description:

  • Added tests for flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put

@vfdev-5 vfdev-5 marked this pull request as ready for review November 19, 2025 21:09
@vfdev-5 vfdev-5 force-pushed the rewrite-prefetch_to_device branch 2 times, most recently from 91e43b4 to 65f7189 Compare December 18, 2025 13:47
@vfdev-5 vfdev-5 marked this pull request as draft March 26, 2026 09:34
@vfdev-5 vfdev-5 force-pushed the rewrite-prefetch_to_device branch from 65f7189 to 920a1d2 Compare April 16, 2026 16:19
@vfdev-5 vfdev-5 changed the title Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put Added tests for flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate Apr 16, 2026
@vfdev-5
Copy link
Copy Markdown
Collaborator Author

vfdev-5 commented Apr 23, 2026

It is fixed on main

@vfdev-5 vfdev-5 closed this Apr 23, 2026
@vfdev-5 vfdev-5 reopened this Apr 23, 2026
@vfdev-5 vfdev-5 force-pushed the rewrite-prefetch_to_device branch from 920a1d2 to b055949 Compare April 23, 2026 21:21
@vfdev-5 vfdev-5 marked this pull request as ready for review April 23, 2026 21:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants