Legacy API (V0)

• Checkpointing with Orbax
  • Saving and Restoring
  • Multiple Objects
  • Managing Checkpoints
  • A Standard Recipe
• API Overview
  • CheckpointManager Layer
    • Basic Usage
    • Managing Multiple Items
  • Checkpointer Layer
    • Saving and Restoring a PyTree
    • Understanding Checkpointers
  • Understanding Items and Registration
  • CheckpointHandler Layer
• Using the Refactored CheckpointManager API
  • Single-Item Checkpointing
    • Before
    • After
  • Multiple-Item Checkpointing
    • Before
    • After
• Checkpointing PyTrees of Arrays
  • CheckpointHandler Support
  • Basic Checkpointing
    • Customizing Restored Properties
    • Partial Restore
  • Advanced Options
    • Saving
    • Restoring
• Checkpoint Format Guide
  • Root Directory
  • Sequential Checkpoint
  • Checkpoint Items
  • PyTree Checkpoints
  • Working with TensorStore
• Optimized Checkpointing with Tensorstore
  • Towards an Improved Checkpoint Format
    • Save Performance (sec)
    • Restore Performance (sec)
  • Checkpoint Format
    • Old Format
    • New Format
  • Enabling the new format
  • Additional Notes
  • Custom Chunk Sizes
  • Customizing Data File Size
    • BEFORE
    • AFTER
• Transformations
• Preemption Tolerance
• Asynchronous Checkpointing
  • What is this?
  • Why should I care?
  • Usage
  • Additional Details
• Debugging Guide
  • Setup
    • Imports
    • Create Sample Checkpoint
  • Checkpoint Size
    • Actual Size on Disk
    • Implied Size from Checkpoint Metadata
  • Tree Metadata
  • Individual Parameters
    • Value Metadata
    • Array Value
• orbax.checkpoint API Reference
  • AbstractCheckpointManager
    • AbstractCheckpointManager
  • CheckpointManager
    • CheckpointManager
    • CheckpointManagerOptions
    • Utility functions
  • Checkpointers
    • AbstractCheckpointer
    • Checkpointer
    • AsyncCheckpointer
    • StandardCheckpointer
    • PyTreeCheckpointer
  • CheckpointHandlers
    • CheckpointHandler
    • AsyncCheckpointHandler
    • StandardCheckpointHandler
    • PyTreeCheckpointHandler
    • CompositeCheckpointHandler
    • JsonCheckpointHandler
    • ArrayCheckpointHandler
    • ProtoCheckpointHandler
    • JaxRandomKeyCheckpointHandler
    • NumpyRandomKeyCheckpointHandler
  • CheckpointArgs
    • CheckpointArgs
    • Composite
    • Standard PyTree
    • Generic PyTree
    • Array
    • JSON
    • Proto
    • JaxRandomKey
    • NumpyRandomKey
    • Utilities
  • TypeHandlers
    • Arguments for PyTreeCheckpointHandler
    • TypeHandler
    • NumpyHandler
    • ScalarHandler
    • ArrayHandler
    • SingleReplicaArrayHandler
    • StringHandler
    • Tensorstore functions
    • TypeHandler registry
  • Atomicity
    • AtomicRenameTemporaryPath
    • CommitFileTemporaryPath
    • Helper functions
  • Atomicity Types
    • TemporaryPath
  • Step entities
    • Metadata
    • NameFormat
    • Factories for various NameFormats
    • Helper functions
    • Helper functions (experimental)
    • Helper functions to implement NameFormat classes
  • Checkpoint Deleter
  • Preservation Policy
    • PreservationContext
    • PreservationPolicy
    • PreserveAll
    • LatestN
    • EveryNSeconds
    • EveryNSteps
    • CustomSteps
    • AnyPreservationPolicy
    • BestN
    • LatestDuration
    • EveryNStepsClosest
  • Save Decision Policy
    • DecisionContext
    • SaveDecisionPolicy
    • FixedIntervalPolicy
    • SpecificStepsPolicy
    • ContinuousCheckpointingPolicy
    • PreemptionCheckpointingPolicy
    • InitialSavePolicy
    • AnySavePolicy
  • Checkpointing Utilities
    • wait_for_new_checkpoint()
    • checkpoints_iterator()
    • construct_restore_args()
  • PyTree Transformations
    • Transform
    • apply_transformations()
    • merge_trees()
  • Msgpack utils
    • msgpack_serialize()
    • msgpack_restore()
  • Multi–host Utilities
    • broadcast_one_to_all()
    • is_primary_host()
    • reached_preemption()
    • sync_global_processes()
    • process_index()
    • BarrierSyncFn
    • get_barrier_sync_fn()
  • Tree Utilities
    • get_param_names()
    • serialize_tree()
    • deserialize_tree()
    • to_flat_dict()
    • from_flat_dict()
    • to_shape_dtype_struct()
  • Metadata Utilities
    • Tree Metadata
    • Sharding Metadata
  • Checkpoint logging
    • AbstractLogger
    • CompositeLogger
    • StandardLogger
    • CloudLogger
    • CloudLoggerOptions
  • Checkpointing Options
    • AsyncOptions
    • MultiprocessingOptions
    • FileOptions
  • Handler Registration
  • Arrays
    • to_shape_dtype_struct()