orbax.checkpoint API Reference

• AbstractCheckpointManager
  • AbstractCheckpointManager
    • AbstractCheckpointManager
• CheckpointManager
  • CheckpointManager
    • CheckpointManager
  • CheckpointManagerOptions
    • CheckpointManagerOptions
  • AsyncOptions
    • AsyncOptions
  • MultiprocessingOptions
    • MultiprocessingOptions
  • Utility functions
    • is_async_checkpointer()
• Checkpointers
  • AbstractCheckpointer
    • AbstractCheckpointer
  • Checkpointer
    • Checkpointer
  • AsyncCheckpointer
    • AsyncCheckpointer
  • StandardCheckpointer
    • StandardCheckpointer
  • PyTreeCheckpointer
    • PyTreeCheckpointer
• CheckpointHandlers
  • CheckpointHandler
    • CheckpointHandler
  • AsyncCheckpointHandler
    • AsyncCheckpointHandler
  • StandardCheckpointHandler
    • StandardCheckpointHandler
  • PyTreeCheckpointHandler
    • PyTreeCheckpointHandler
  • CompositeCheckpointHandler
    • CompositeCheckpointHandler
  • JsonCheckpointHandler
    • JsonCheckpointHandler
  • ArrayCheckpointHandler
    • ArrayCheckpointHandler
  • ProtoCheckpointHandler
    • ProtoCheckpointHandler
  • JaxRandomKeyCheckpointHandler
    • JaxRandomKeyCheckpointHandler
  • NumpyRandomKeyCheckpointHandler
    • NumpyRandomKeyCheckpointHandler
• CheckpointArgs
  • CheckpointArgs
    • CheckpointArgs
    • CheckpointArgs
  • Array
    • ArraySave
    • ArrayRestore
  • Standard PyTree
    • StandardSave
    • StandardRestore
  • Generic PyTree
    • PyTreeSave
    • PyTreeRestore
  • Composite
    • Composite
  • 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
  • SingleSliceArrayHandler
    • ArrayHandler
  • 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()
• 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()
    • is_gcs_path()
    • get_save_directory()
    • is_tmp_checkpoint()
    • is_checkpoint_finalized()
    • create_tmp_directory()
    • tmp_checkpoints()
    • cleanup_tmp_directories()
    • get_tmp_directory()
    • step_from_checkpoint_name()
  • Helper functions to implement NameFormat classes
    • build_step_metadatas()
    • step_prefix_with_underscore()
• Checkpointing Utilities
  • wait_for_new_checkpoint()
  • unlock_existing_checkpoints()
  • checkpoints_iterator()
  • construct_restore_args()
• General Utilities
  • Async wrappers
    • async_makedirs()
    • async_write_bytes()
    • async_exists()
  • Tree utils
    • is_empty_or_leaf()
    • get_key_name()
    • to_flat_dict()
    • serialize_tree()
    • deserialize_tree()
    • from_flat_dict()
    • pytree_structure()
  • Aggregate file
    • leaf_is_placeholder()
    • leaf_placeholder()
    • name_from_leaf_placeholder()
    • is_supported_empty_aggregation_type()
    • is_supported_aggregation_type()
  • Directories
    • cleanup_tmp_directories()
    • get_tmp_directory()
    • create_tmp_directory()
    • get_save_directory()
    • is_gcs_path()
  • Atomicity
    • is_tmp_checkpoint()
    • is_checkpoint_finalized()
  • Checkpoint steps
    • step_from_checkpoint_name()
    • checkpoint_steps_paths()
    • checkpoint_steps()
    • any_checkpoint_step()
    • tmp_checkpoints()
    • lockdir()
    • is_locked()
    • are_locked()
  • Sharding
    • fully_replicated_host_local_array_to_global_array()
  • Misc.
    • is_scalar()
    • record_saved_duration()
• PyTree Transformations
  • Transform
    • Transform.__eq__()
    • Transform.__hash__
    • Transform.__init__()
  • apply_transformations()
  • merge_trees()
• Msgpack utils
  • msgpack_serialize()
  • msgpack_restore()
• Multi-host Utilities
  • sync_global_processes()
  • broadcast_one_to_all()
  • broadcast_one_to_some()
• Metadata Utilities
  • Tree Metadata
    • Metadata
    • ArrayMetadata
    • ScalarMetadata
    • StringMetadata
  • Sharding Metadata
    • ShardingMetadata
    • NamedShardingMetadata
    • SingleDeviceShardingMetadata
    • GSPMDShardingMetadata
    • PositionalShardingMetadata
    • ShardingTypes
    • from_jax_sharding()
    • from_serialized_string()
    • get_sharding_or_none()
  • Internal Metadata
    • TreeMetadata
• Checkpoint logging
  • AbstractLogger
  • CompositeLogger
  • StandardLogger
  • CloudLogger
  • CloudLoggerOptions
  • StepStatistics