orbax.checkpoint API Reference

• AbstractCheckpointManager
  • AbstractCheckpointManager
    • AbstractCheckpointManager
• CheckpointManager
  • CheckpointManager
    • CheckpointManager
  • CheckpointManagerOptions
    • CheckpointManagerOptions
  • Utility functions
    • is_async_checkpointer()
• Checkpointers
  • AbstractCheckpointer
    • AbstractCheckpointer
  • Checkpointer
    • Checkpointer
  • AsyncCheckpointer
    • AsyncCheckpointer
  • StandardCheckpointer
    • StandardCheckpointer
  • PyTreeCheckpointer
    • PyTreeCheckpointer
• CheckpointHandlers
  • CheckpointHandler
    • CheckpointHandler
  • AsyncCheckpointHandler
    • AsyncCheckpointHandler
  • StandardCheckpointHandler
    • StandardCheckpointHandler
    • StandardSaveArgs
    • StandardRestoreArgs
  • PyTreeCheckpointHandler
    • PyTreeCheckpointHandler
    • PyTreeSaveArgs
    • PyTreeRestoreArgs
  • CompositeCheckpointHandler
    • CompositeCheckpointHandler
    • CompositeArgs
    • CompositeResults
  • JsonCheckpointHandler
    • JsonCheckpointHandler
    • JsonSaveArgs
    • JsonRestoreArgs
  • ArrayCheckpointHandler
    • ArrayCheckpointHandler
    • ArraySaveArgs
    • ArrayRestoreArgs
  • ProtoCheckpointHandler
    • ProtoCheckpointHandler
    • ProtoSaveArgs
    • ProtoRestoreArgs
  • JaxRandomKeyCheckpointHandler
    • JaxRandomKeyCheckpointHandler
    • JaxRandomKeySaveArgs
    • JaxRandomKeyRestoreArgs
  • NumpyRandomKeyCheckpointHandler
    • NumpyRandomKeyCheckpointHandler
    • NumpyRandomKeySaveArgs
    • NumpyRandomKeyRestoreArgs
• CheckpointArgs
  • CheckpointArgs
    • CheckpointArgs
  • Composite
    • Composite
  • Standard PyTree
    • StandardSave
    • StandardRestore
  • Generic PyTree
    • PyTreeSave
    • PyTreeRestore
  • Array
    • ArraySave
    • ArrayRestore
  • JSON
    • JsonSave
    • JsonRestore
  • Proto
    • ProtoSave
    • ProtoRestore
  • JaxRandomKey
    • JaxRandomKeySave
    • JaxRandomKeyRestore
  • NumpyRandomKey
    • NumpyRandomKeySave
    • NumpyRandomKeyRestore
  • Utilities
    • get_registered_handler_cls()
    • register_with_handler()
    • get_registered_args_cls()
    • has_registered_args()
• TypeHandlers
  • Arguments for PyTreeCheckpointHandler
    • SaveArgs
    • RestoreArgs
    • ArrayRestoreArgs
  • TypeHandler
    • TypeHandler
  • NumpyHandler
    • NumpyHandler
  • ScalarHandler
    • ScalarHandler
  • ArrayHandler
    • ArrayHandler
  • SingleReplicaArrayHandler
    • SingleReplicaArrayHandler
  • StringHandler
    • StringHandler
  • Tensorstore functions
    • is_ocdbt_checkpoint()
    • merge_ocdbt_per_process_files()
    • get_json_tspec_write()
    • get_json_tspec_read()
    • get_ts_context()
    • get_cast_tspec_serialize()
    • get_cast_tspec_deserialize()
  • TypeHandler registry
    • TypeHandlerRegistry
    • create_type_handler_registry()
    • register_type_handler()
    • get_type_handler()
    • has_type_handler()
    • register_standard_handlers_with_options()
• Atomicity
  • AtomicRenameTemporaryPath
    • AtomicRenameTemporaryPath
  • CommitFileTemporaryPath
    • CommitFileTemporaryPath
  • Helper functions
    • on_commit_callback()
• Atomicity Types
  • TemporaryPath
    • TemporaryPath
• Step entities
  • Metadata
    • Metadata
  • NameFormat
    • NameFormat
  • Factories for various NameFormats
    • standard_name_format()
    • composite_name_format()
  • Helper functions
    • build_step_path()
    • checkpoint_steps_paths()
    • checkpoint_steps()
    • any_checkpoint_step()
    • latest_step_metadata()
    • step_metadata_of_checkpoint_path()
  • Helper functions (experimental)
    • find_step_path()
    • get_save_directory()
    • is_tmp_checkpoint()
    • is_checkpoint_finalized()
    • tmp_checkpoints()
    • step_from_checkpoint_name()
  • Helper functions to implement NameFormat classes
    • build_step_metadatas()
    • step_prefix_with_underscore()
• Checkpoint Deleter
• Preservation Policy
  • PreservationContext
    • PreservationContext
  • PreservationPolicy
    • PreservationPolicy
  • PreserveAll
    • PreserveAll
  • LatestN
    • LatestN
  • EveryNSeconds
    • EveryNSeconds
  • EveryNSteps
    • EveryNSteps
  • CustomSteps
    • CustomSteps
  • AnyPreservationPolicy
    • AnyPreservationPolicy
  • BestN
    • BestN
• Save Decision Policy
  • DecisionContext
    • DecisionContext
  • SaveDecisionPolicy
    • SaveDecisionPolicy
  • FixedIntervalPolicy
    • FixedIntervalPolicy
  • SpecificStepsPolicy
    • SpecificStepsPolicy
  • ContinuousCheckpointingPolicy
    • ContinuousCheckpointingPolicy
  • PreemptionCheckpointingPolicy
    • PreemptionCheckpointingPolicy
  • InitialSavePolicy
    • InitialSavePolicy
  • AnySavePolicy
    • AnySavePolicy
• Checkpointing Utilities
  • wait_for_new_checkpoint()
  • checkpoints_iterator()
  • construct_restore_args()
• PyTree Transformations
  • Transform
    • Transform.__eq__()
    • Transform.__hash__
    • Transform.__init__()
  • 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
    • BarrierSyncFn.__call__()
    • BarrierSyncFn.__init__()
    • BarrierSyncFn.__subclasshook__()
  • 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
    • ArrayMetadata
    • ScalarMetadata
    • StringMetadata
    • StorageMetadata
  • Sharding Metadata
    • ShardingMetadata
    • NamedShardingMetadata
    • SingleDeviceShardingMetadata
• Checkpoint logging
  • AbstractLogger
    • AbstractLogger
  • CompositeLogger
    • CompositeLogger
  • StandardLogger
    • StandardLogger
  • CloudLogger
    • CloudLogger
  • CloudLoggerOptions
    • CloudLoggerOptions
• Checkpointing Options
  • AsyncOptions
    • AsyncOptions
  • MultiprocessingOptions
    • MultiprocessingOptions
  • FileOptions
    • FileOptions
• Handler Registration
• Arrays
  • to_shape_dtype_struct()