orbax
orbax copied to clipboard
google
Orbax provides common utility libraries for JAX users.
Use `ocp.future.CommitFutureAwaitingContractedSignals` in custom async checkpoint handlers and type handlers to support async directory creation.
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
Patch for testing b/404645078, hardcoding single_host_load_and_broadcast=true.
Add `protobuf` to Orbax Export dependencies and install protobuf using python and ensure `protoc` installation in github workflow.
I am getting a number of errors checking for folder existence in _src/metadata/checkpoint.py e.g. ```_src/metadata/checkpoint.py", line 45, in _sanitize_metadata_path raise FileNotFoundError(f'Path does not exist: {path}')``` when trying to create a...
Fix the args ordering when using assertProto2Equal for OBM related tests
Allow configuring the minimum interval seconds for continuous checkpointing.
Simplify heartbeat options for PjRt distributed runtime.