orbax

orbax copied to clipboard

#v1 Add `LeafHandler` as a `CheckpointableHandler`, so that ordinary PyTree leaves can also be saved as individual checkpointables.

copybara-service[bot]

ImportError: cannot import name 'oex_orchestration_pb2' from 'orbax.export.protos'

3

When running the latest `orbax-export`, the `from orbax.export.protos import oex_orchestration_pb2` line in utils.py is failing with an `ImportError: cannot import name 'oex_orchestration_pb2' from 'orbax.export.protos'`. Here is a test log: https://github.com/jax-ml/jax-ai-stack/actions/runs/17402977130/job/49400221219.

kanglant

Internal change.

copybara-service[bot]

Internal adjustment to tree-verity reporting

1

Internal adjustment to tree-verity reporting

copybara-service[bot]

Migrate Tunelab to Orbax `PreservationPolicy` for checkpoint management.

copybara-service[bot]

Fix _glob_step_paths if HNS is enabled

1

Seems like there was an oversight in the case where HNS is enabled, but no `step_prefix` is provided. In that case, `step_prefix` is None, which results in an error when...

dvruette

`manager.restore` now returns a `orbax.checkpoint._src.handlers.composite_checkpoint_handler` object instead of a PyTree, so we need to access the `state` member of the returned object.

keshavb96

Internal change

copybara-service[bot]

Internal change

copybara-service[bot]

Update sidv2, yt_prod_v3 test data and re-enable the test.

copybara-service[bot]