orbax
orbax copied to clipboard
google
Orbax provides common utility libraries for JAX users.
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
Add user-facing documentation for checkpoint handler registry.
Add a registry for `CheckpointHandler`s.
Hi Everyone, I've been trying to checkpoint training using Orbax in a project linked here [project](https://github.com/ssenan/dnadiffusion-jax/tree/main). When I test the code locally I'm able checkpoint successfully, but when training in...
After upgrading to `jax==0.4.31` I am seeing this error when trying to save a model using the `PyTreeCheckpointer`. Downgrading to `0.4.30` fixed it for now. ``` File ".../site-packages/orbax/checkpoint/checkpointer.py", line 151,...
Include directory path in `create_directory` barrier key.