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

Open

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

copybara-service wants to merge 1 commit into main from test_846723995

copybara-service

danielsuo

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

3f5337ca

Login to write a write a comment.

Login via GitHub

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone