orbax
orbax copied to clipboard
google
Orbax provides common utility libraries for JAX users.
I was trying to find info about how to export an orbax checkpoint to SavedModel, and found broken links or code examples in the docs with exception / stack traces:...
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
Support lazy init of a CheckpointInfo with `CheckpointManagerOptions.lazy_checkpoint_info`.
Add `tree_trim` function for filtering a PyTree by another's structure.
Catching very infrequent error in train-logs writing which crashes train job.
Patch for `use_replica_parallel=False`.
Use ocp.futures.CommitFutureAwaitingContractedSignals to support async directory creation feature for Gemax Prod
The example code from the documentation of the form ``` python mngr.save(0, args=ocp.args.PyTreeSave(tree)) ``` shows a type error saying "expected 0 positional arguments" (to PyTreeSave) in vscodium with pyright. Playing...
