jax
d962a2bb - [Pallas] Use NumPy CPU random number generators instead of generating on TPU.

Commit

69 days ago

[Pallas] Use NumPy CPU random number generators instead of generating on TPU. It seems that such loop fusion for huge shapes are slow under asan. PiperOrigin-RevId: 872423458

References

#35245 - [Pallas] Use NumPy CPU random number generators instead of generating on TPU.

Author

WindQAQ

Committer

Google-ML-Automation

Parents

115ceae8

jax d962a2bb - [Pallas] Use NumPy CPU random number generators instead of generating on TPU.

jax
d962a2bb - [Pallas] Use NumPy CPU random number generators instead of generating on TPU.