flax
[pmap] Avoid degraded performance under the new `jax.pmap`.
#5152

Merged

[pmap] Avoid degraded performance under the new `jax.pmap`. #5152

copybara-service merged 1 commit into main from test_846723995

copybara-service force pushed from 3f5337ca to b18b3968 172 days ago

[pmap] Avoid degraded performance under the new `jax.pmap`.

5d908519

copybara-service force pushed from b18b3968 to 5d908519 172 days ago

copybara-service merged 5d908519 into main 172 days ago

copybara-service deleted the test_846723995 branch 172 days ago

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone