jax
8777357b - [Mosaic TPU] Fix up some broken tests

Commit

323 days ago

[Mosaic TPU] Fix up some broken tests Our recent changes made us more strict in our VMEM allocations, exposing some OOMs. PiperOrigin-RevId: 827995654

References

#33091 - [Mosaic TPU] Fix up some broken tests

Author

apaszke

Committer

Google-ML-Automation

Parents

9bad30c1

jax 8777357b - [Mosaic TPU] Fix up some broken tests

jax
8777357b - [Mosaic TPU] Fix up some broken tests