Orbax: Distributed Checkpointing with JAX

Colin Gaffney, Shutong Li, Daniel Ng +13

In a landscape of high-performance distributed ML systems, JAX has emerged as a framework of choice. However, JAX's modular design philosophy leaves it without a standardized check…