Emergency checkpoint: use JAX for global_max and combine multiple broadcasts

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users

https://orbax.readthedocs.io/

Apache License 2.0

304 stars 36 forks source link

Closed copybara-service[bot] closed 2 weeks ago

copybara-service[bot] commented 2 weeks ago

Emergency checkpoint: use JAX for global_max and combine multiple broadcasts