Skip to content

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

Merged
copybara-service[bot] merged 1 commit into
mainfrom
rewrite-prefetch_to_device
Apr 28, 2026
Merged

copybara-service[bot] merged 1 commit into
mainfrom
rewrite-prefetch_to_device

Conversation

@vfdev-5

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

Copy link
Copy Markdown
Collaborator

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

vfdev-5 commented Apr 23, 2026 •

Copy link
Copy Markdown
Collaborator Author

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
@copybara-service
copybara-service Bot merged commit c781023 into main Apr 28, 2026
24 checks passed
@copybara-service
copybara-service Bot deleted the rewrite-prefetch_to_device branch April 28, 2026 22:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants