orbax
orbax copied to clipboard
google
Orbax provides common utility libraries for JAX users.
Add benchmarks comparing with/without asynchronous directory creation
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
Allow users to enable continuous checkpointing in Gemax Prod. Note that this is an experimental feature and should not be adopted by users without explicit approval.
Add support for merging array groups across multiple devices, with uneven sharding between devices.
Remove items_and_handlers-related `CompositeCheckpointHandler` code.
Resubmit of https://github.com/google/orbax/pull/1304 with style adjustments.
