jax
23d5c10f - [Mosaic:TPU] Fix fully replicated relayout

Commit

1 year ago

[Mosaic:TPU] Fix fully replicated relayout It was incorrect since batch dims are not replicated PiperOrigin-RevId: 703189919

References

Author

tlongeri

Committer

Google-ML-Automation

Parents