orbax
orbax copied to clipboard
Orbax provides common utility libraries for JAX users.
Add benchmarks comparing with/without asynchronous directory creation
Add pytree_checkpoint_handler_test.
Add checkpoint/proto/testing/BUILD.
Internal
Internal
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.
Refactor metadata de/serialization.
Internal
Resubmit of https://github.com/google/orbax/pull/1304 with style adjustments.