Orbax: Distributed Checkpointing with JAX. - Colin Gaffney; Shutong Li; Daniel Ng; Anastasia Petrushkina; Niket Kumar; Adam Cogdell; Mridul Sahu; Yaning Liang; Nikhil Bansal; Justin Pan; Angel Mau; Abhishek Agrawal; et al. (CoRR, 2026)