Skip to content

v0.12.0

Choose a tag to compare

@orbax-dev orbax-dev released this 02 Jun 20:46

Added

  • Add LatestDuration preservation policy to support duration-based (TTL)
    checkpoint retention.
  • Add EveryNStepsClosest preservation policy, which preserves checkpoints at
    steps closest to multiples of interval_steps.
  • Add SerializationStatusCallback for save operations to track transfer and
    disk write events. While primarily intended as lightweight progress handles,
    this callback system allows the potential to inject custom logic for
    advanced monitoring and custom actions during checkpointing.

Changed

  • #v1 Rename save/load_pytree to save/load. Eliminate most user-facing
    "pytree" terminology in favor of "state" as a more specific term.
    Add deprecations.py for handling deprecated public functions.

Fixed

  • Fix Pathways MTC Colocated worker initialization and save/restore PyTree bugs.
  • Fix bug where TensorStore transaction was never used, due to byte_limiter
    always having a non-None value.