orbax
orbax copied to clipboard
google
Orbax provides common utility libraries for JAX users.
Add SaveDecisionPolicy and PreservationPolicy to v1, use CheckpointMetadata instead of CheckpointInfo.
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
Add Checkpointable interface and GenericHandler, and implement save & load methods within GenericHandler.
#v1 Add checkpointables support for `training.Checkpointer`.
Don't raise errors when directories don't exist. Context: GCS doesn't support natively empty dirs. When using etils[epath-gcs] without tensorflow, but with the gcsfs backend, path.mkdirs does not create empty dirs....
Rename `best_fn` to `get_metric_fn` in CheckpointManagerOptions.
#v1 Add documentation on `checkpointables`. This deals primarily with the concept of what a checkpointable is, as well as usage examples.
Updated save and restore signatures for Emergency Checkpoint Manager
#v1 Add compatibility tests for save-by-v0-load-by-v1 and also fix code.
Hi, I'm trying to use orbax to checkpoint arbitrary penzai models, which have parameters of type `pz.ParameterValue`, which themselves contain `pz.NamedArray` instances. These instances contain named axes. So, I thought...
