flax
[graphcast] Prepare for `jax_pmap_shmap_merge=True`.
#5168

Open

[graphcast] Prepare for `jax_pmap_shmap_merge=True`. #5168

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

copybara-service

danielsuo

[graphcast] Prepare for `jax_pmap_shmap_merge=True`.

e3ebf2fa

Login to write a write a comment.

Login via GitHub

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone