jax
d8ac2546 - Fix remaining tests under `pmap_shmap_merge=True`.

Commit

129 days ago

Fix remaining tests under `pmap_shmap_merge=True`. PiperOrigin-RevId: 810556321

References

#31917 - Fix remaining tests under `pmap_shmap_merge=True`.

Author

danielsuo

danielsuo

Committer

Google-ML-Automation

Google-ML-Automation

Parents

Loading