Index

_

__delattr__() (orbax.checkpoint.path.step.Metadata method) __eq__() (orbax.checkpoint.args.CheckpointArgs method), [1] (orbax.checkpoint.checkpoint_manager.AsyncOptions method) (orbax.checkpoint.checkpoint_manager.CheckpointManagerOptions method) (orbax.checkpoint.checkpoint_manager.MultiprocessingOptions method) (orbax.checkpoint.metadata.ArrayMetadata method) (orbax.checkpoint.metadata.GSPMDShardingMetadata method) (orbax.checkpoint.metadata.Metadata method) (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.PositionalShardingMetadata method) (orbax.checkpoint.metadata.ScalarMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method) (orbax.checkpoint.metadata.StringMetadata method) (orbax.checkpoint.metadata.TreeMetadata method) (orbax.checkpoint.path.step.Metadata method) (orbax.checkpoint.transform_utils.Transform method) (orbax.checkpoint.type_handlers.ArrayRestoreArgs method) (orbax.checkpoint.type_handlers.RestoreArgs method) (orbax.checkpoint.type_handlers.SaveArgs method) (orbax.export.serving_config.ServingConfig method) (orbax.export.utils.TensorSpecWithDefault method) (orbax.export.validate.validation_report.ValidationReportOption method) __hash__ (orbax.checkpoint.args.CheckpointArgs attribute), [1] (orbax.checkpoint.checkpoint_manager.AsyncOptions attribute) (orbax.checkpoint.checkpoint_manager.CheckpointManagerOptions attribute) (orbax.checkpoint.checkpoint_manager.MultiprocessingOptions attribute) (orbax.checkpoint.metadata.ArrayMetadata attribute) (orbax.checkpoint.metadata.GSPMDShardingMetadata attribute) (orbax.checkpoint.metadata.Metadata attribute) (orbax.checkpoint.metadata.NamedShardingMetadata attribute) (orbax.checkpoint.metadata.PositionalShardingMetadata attribute) (orbax.checkpoint.metadata.ScalarMetadata attribute) (orbax.checkpoint.metadata.ShardingMetadata attribute) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata attribute) (orbax.checkpoint.metadata.StringMetadata attribute) (orbax.checkpoint.metadata.TreeMetadata attribute) (orbax.checkpoint.transform_utils.Transform attribute) (orbax.checkpoint.type_handlers.ArrayRestoreArgs attribute) (orbax.checkpoint.type_handlers.RestoreArgs attribute) (orbax.checkpoint.type_handlers.SaveArgs attribute) (orbax.export.serving_config.ServingConfig attribute) (orbax.export.utils.TensorSpecWithDefault attribute) (orbax.export.validate.validation_report.ValidationReportOption attribute) __hash__() (orbax.checkpoint.path.step.Metadata method) __init__() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.args.CheckpointArgs method), [1] (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.checkpoint_manager.AsyncOptions method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManagerOptions method) (orbax.checkpoint.checkpoint_manager.MultiprocessingOptions method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.JsonCheckpointHandler method) (orbax.checkpoint.metadata.ArrayMetadata method) (orbax.checkpoint.metadata.GSPMDShardingMetadata method) (orbax.checkpoint.metadata.Metadata method) (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.PositionalShardingMetadata method) (orbax.checkpoint.metadata.ScalarMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method) (orbax.checkpoint.metadata.StringMetadata method) (orbax.checkpoint.metadata.TreeMetadata method) (orbax.checkpoint.path.step.Metadata method) (orbax.checkpoint.path.step.NameFormat method) (orbax.checkpoint.ProtoCheckpointHandler method) (orbax.checkpoint.PyTreeCheckpointer method) (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointer method) (orbax.checkpoint.StandardCheckpointHandler method) (orbax.checkpoint.transform_utils.Transform method) (orbax.checkpoint.type_handlers.ArrayHandler method), [1] (orbax.checkpoint.type_handlers.ArrayRestoreArgs method) (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.RestoreArgs method) (orbax.checkpoint.type_handlers.SaveArgs method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.export.export_manager.ExportManager method) (orbax.export.jax_module.JaxModule method) (orbax.export.serving_config.ServingConfig method) (orbax.export.utils.CallableSignatures method) (orbax.export.utils.TensorSpecWithDefault method) (orbax.export.validate.validation_manager.ValidationManager method) (orbax.export.validate.validation_report.ValidationReport method) (orbax.export.validate.validation_report.ValidationReportOption method)

__post_init__() (orbax.export.validate.validation_report.ValidationReportOption method) __setattr__() (orbax.checkpoint.path.step.Metadata method) __subclasshook__() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) (orbax.checkpoint.path.step.NameFormat method)

A

AbstractCheckpointer (class in orbax.checkpoint) AbstractCheckpointManager (class in orbax.checkpoint.abstract_checkpoint_manager) all_steps() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) any_checkpoint_step() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) apply_transformations() (in module orbax.checkpoint.transform_utils) are_locked() (in module orbax.checkpoint.utils) ArrayCheckpointHandler (class in orbax.checkpoint) ArrayHandler (class in orbax.checkpoint.type_handlers), [1] ArrayMetadata (class in orbax.checkpoint.metadata) ArrayRestore (in module orbax.checkpoint.args) ArrayRestoreArgs (class in orbax.checkpoint.type_handlers) ArraySave (in module orbax.checkpoint.args)

as_nested_tree() (orbax.checkpoint.metadata.TreeMetadata method) async_exists() (in module orbax.checkpoint.utils) async_makedirs() (in module orbax.checkpoint.utils) async_save() (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointHandler method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.JaxRandomKeyCheckpointHandler method) (orbax.checkpoint.NumpyRandomKeyCheckpointHandler method) (orbax.checkpoint.ProtoCheckpointHandler method) (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) async_write_bytes() (in module orbax.checkpoint.utils) AsyncCheckpointer (class in orbax.checkpoint) AsyncCheckpointHandler (class in orbax.checkpoint) AsyncOptions (class in orbax.checkpoint.checkpoint_manager)

B

best_step() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) bind() (orbax.export.serving_config.ServingConfig method) broadcast_one_to_all() (in module orbax.checkpoint.multihost) broadcast_one_to_some() (in module orbax.checkpoint.multihost)

build() (orbax.checkpoint.metadata.TreeMetadata class method) build_metadata() (orbax.checkpoint.path.step.NameFormat method) build_name() (orbax.checkpoint.path.step.NameFormat method) build_step_metadatas() (in module orbax.checkpoint.path.step) build_step_path() (in module orbax.checkpoint.path.step)

C

CallableSignatures (class in orbax.export.utils) check_for_errors() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) check_input() (orbax.export.validate.validation_manager.ValidationManager class method) check_output() (orbax.export.validate.validation_manager.ValidationManager class method) checkpoint_steps() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) checkpoint_steps_paths() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) CheckpointArgs (class in orbax.checkpoint.args), [1] Checkpointer (class in orbax.checkpoint) CheckpointHandler (class in orbax.checkpoint) CheckpointManager (class in orbax.checkpoint.checkpoint_manager) CheckpointManagerOptions (class in orbax.checkpoint.checkpoint_manager) checkpoints_iterator() (in module orbax.checkpoint.checkpoint_utils) cleanup_tmp_directories() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils)

close() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.ProtoCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) commit_timestamp (orbax.checkpoint.path.step.Metadata property) commit_timestamp_nsecs (orbax.checkpoint.path.step.Metadata property) Composite (in module orbax.checkpoint.args) composite_name_format() (in module orbax.checkpoint.path.step) CompositeCheckpointHandler (class in orbax.checkpoint) construct_restore_args() (in module orbax.checkpoint.checkpoint_utils) create_tmp_directory() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) create_type_handler_registry() (in module orbax.checkpoint.type_handlers)

D

delete() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) deserialize() (orbax.checkpoint.type_handlers.ArrayHandler method), [1] (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.ScalarHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.type_handlers.TypeHandler method)

deserialize_tree() (in module orbax.checkpoint.utils) directory (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager property) (orbax.checkpoint.checkpoint_manager.CheckpointManager property) dtensor_initialized() (in module orbax.export.dtensor_utils)

E

ExportManager (class in orbax.export.export_manager)

ExportManagerBase (class in orbax.export.export_manager_base)

F

finalize() (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) (orbax.checkpoint.type_handlers.TypeHandler method) find_all() (orbax.checkpoint.path.step.NameFormat method) find_metadata() (orbax.checkpoint.path.step.NameFormat method) find_step() (orbax.checkpoint.path.step.NameFormat method) find_step_path() (in module orbax.checkpoint.path.step) from_deserialized_dict() (orbax.checkpoint.metadata.NamedShardingMetadata class method) (orbax.checkpoint.metadata.ShardingMetadata class method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata class method)

from_flat_dict() (in module orbax.checkpoint.utils) from_jax_sharding() (in module orbax.checkpoint.metadata) (orbax.checkpoint.metadata.NamedShardingMetadata class method) (orbax.checkpoint.metadata.ShardingMetadata class method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata class method) from_json() (orbax.checkpoint.metadata.TreeMetadata class method) from_saved_model() (orbax.export.utils.CallableSignatures class method) from_serialized_string() (in module orbax.checkpoint.metadata) fully_replicated_host_local_array_to_global_array() (in module orbax.checkpoint.utils)

G

get_cast_tspec_deserialize() (in module orbax.checkpoint.type_handlers) get_cast_tspec_serialize() (in module orbax.checkpoint.type_handlers) get_current_dtensor_mesh() (in module orbax.export.dtensor_utils) get_infer_step() (orbax.export.serving_config.ServingConfig method) get_input_signature() (orbax.export.serving_config.ServingConfig method) get_json_tspec_read() (in module orbax.checkpoint.type_handlers) get_json_tspec_write() (in module orbax.checkpoint.type_handlers) get_key_name() (in module orbax.checkpoint.utils) get_registered_args_cls() (in module orbax.checkpoint.args)

get_registered_handler_cls() (in module orbax.checkpoint.args) get_save_directory() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) get_sharding_or_none() (in module orbax.checkpoint.metadata) get_tmp_directory() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) get_ts_context() (in module orbax.checkpoint.type_handlers) get_type_handler() (in module orbax.checkpoint.type_handlers) GSPMDShardingMetadata (class in orbax.checkpoint.metadata)

H

has_registered_args() (in module orbax.checkpoint.args)

has_type_handler() (in module orbax.checkpoint.type_handlers)

I

init_timestamp_nsecs (orbax.checkpoint.path.step.Metadata property) initialize_dtensor() (in module orbax.export.dtensor_utils) is_async_checkpointer() (in module orbax.checkpoint.checkpoint_manager) is_checkpoint_finalized() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) is_empty_or_leaf() (in module orbax.checkpoint.utils) is_gcs_path() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils)

is_locked() (in module orbax.checkpoint.utils) is_ocdbt_checkpoint() (in module orbax.checkpoint.type_handlers) is_scalar() (in module orbax.checkpoint.utils) is_supported_aggregation_type() (in module orbax.checkpoint.utils) is_supported_empty_aggregation_type() (in module orbax.checkpoint.utils) is_tmp_checkpoint() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) item_metadata() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method)

J

jax_mesh_to_dtensor_mesh() (in module orbax.export.dtensor_utils) jax_methods (orbax.export.jax_module.JaxModule property) JaxModule (class in orbax.export.jax_module) JaxRandomKeyCheckpointHandler (class in orbax.checkpoint)

JaxRandomKeyRestore (in module orbax.checkpoint.args) JaxRandomKeySave (in module orbax.checkpoint.args) JsonCheckpointHandler (class in orbax.checkpoint) JsonRestore (in module orbax.checkpoint.args) JsonSave (in module orbax.checkpoint.args)

L

latest_step() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) latest_step_metadata() (in module orbax.checkpoint.path.step) leaf_is_placeholder() (in module orbax.checkpoint.utils)

leaf_placeholder() (in module orbax.checkpoint.utils) load() (orbax.export.export_manager.ExportManager method) (orbax.export.export_manager_base.ExportManagerBase method) lockdir() (in module orbax.checkpoint.utils)

M

make_auto_batching_function() (in module orbax.export.utils) maybe_enable_dtensor_export_on() (in module orbax.export.dtensor_utils) merge_ocdbt_per_process_files() (in module orbax.checkpoint.type_handlers) merge_trees() (in module orbax.checkpoint.transform_utils) Metadata (class in orbax.checkpoint.metadata) (class in orbax.checkpoint.path.step) metadata() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) (orbax.checkpoint.type_handlers.ArrayHandler method), [1] (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.ScalarHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.type_handlers.TypeHandler method) methods (orbax.export.jax_module.JaxModule property) metrics() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) module orbax.checkpoint, [1] orbax.checkpoint.abstract_checkpoint_manager orbax.checkpoint.args orbax.checkpoint.checkpoint_manager orbax.checkpoint.checkpoint_utils orbax.checkpoint.logging orbax.checkpoint.metadata orbax.checkpoint.msgpack_utils orbax.checkpoint.multihost orbax.checkpoint.path.step orbax.checkpoint.transform_utils orbax.checkpoint.type_handlers orbax.checkpoint.utils orbax.export.dtensor_utils orbax.export.export_manager orbax.export.export_manager_base orbax.export.jax_module orbax.export.serving_config orbax.export.utils orbax.export.validate.validation_manager orbax.export.validate.validation_report

msgpack_restore() (in module orbax.checkpoint.msgpack_utils) msgpack_serialize() (in module orbax.checkpoint.msgpack_utils) MultiprocessingOptions (class in orbax.checkpoint.checkpoint_manager)

N

name_from_leaf_placeholder() (in module orbax.checkpoint.utils) NamedShardingMetadata (class in orbax.checkpoint.metadata) NameFormat (class in orbax.checkpoint.path.step)

NumpyHandler (class in orbax.checkpoint.type_handlers) NumpyRandomKeyCheckpointHandler (class in orbax.checkpoint) NumpyRandomKeyRestore (in module orbax.checkpoint.args) NumpyRandomKeySave (in module orbax.checkpoint.args)

O

orbax.checkpoint module, [1] orbax.checkpoint.abstract_checkpoint_manager module orbax.checkpoint.args module orbax.checkpoint.checkpoint_manager module orbax.checkpoint.checkpoint_utils module orbax.checkpoint.logging module orbax.checkpoint.metadata module orbax.checkpoint.msgpack_utils module orbax.checkpoint.multihost module orbax.checkpoint.path.step module orbax.checkpoint.transform_utils module

orbax.checkpoint.type_handlers module orbax.checkpoint.utils module orbax.export.dtensor_utils module orbax.export.export_manager module orbax.export.export_manager_base module orbax.export.jax_module module orbax.export.serving_config module orbax.export.utils module orbax.export.validate.validation_manager module orbax.export.validate.validation_report module

P

path (orbax.checkpoint.path.step.Metadata attribute) PositionalShardingMetadata (class in orbax.checkpoint.metadata) ProtoCheckpointHandler (class in orbax.checkpoint) ProtoRestore (in module orbax.checkpoint.args) ProtoSave (in module orbax.checkpoint.args)

pytree_structure() (in module orbax.checkpoint.utils) PyTreeCheckpointer (class in orbax.checkpoint) PyTreeCheckpointHandler (class in orbax.checkpoint) PyTreeRestore (in module orbax.checkpoint.args) PyTreeSave (in module orbax.checkpoint.args)

R

reached_preemption() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) record_saved_duration() (in module orbax.checkpoint.utils) register_standard_handlers_with_options() (in module orbax.checkpoint.type_handlers) register_type_handler() (in module orbax.checkpoint.type_handlers) register_with_handler() (in module orbax.checkpoint.args) reload() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) restore() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.JaxRandomKeyCheckpointHandler method) (orbax.checkpoint.JsonCheckpointHandler method) (orbax.checkpoint.NumpyRandomKeyCheckpointHandler method) (orbax.checkpoint.ProtoCheckpointHandler method) (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method)

restore_type (orbax.checkpoint.type_handlers.ArrayRestoreArgs attribute) RestoreArgs (class in orbax.checkpoint.type_handlers)

S

save() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.JaxRandomKeyCheckpointHandler method) (orbax.checkpoint.JsonCheckpointHandler method) (orbax.checkpoint.NumpyRandomKeyCheckpointHandler method) (orbax.checkpoint.ProtoCheckpointHandler method) (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) (orbax.export.export_manager.ExportManager method) (orbax.export.export_manager_base.ExportManagerBase method) SaveArgs (class in orbax.checkpoint.type_handlers) ScalarHandler (class in orbax.checkpoint.type_handlers) ScalarMetadata (class in orbax.checkpoint.metadata) serialize() (orbax.checkpoint.type_handlers.ArrayHandler method), [1] (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.ScalarHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.type_handlers.TypeHandler method)

serialize_tree() (in module orbax.checkpoint.utils) serving_signatures (orbax.export.export_manager.ExportManager property) ServingConfig (class in orbax.export.serving_config) ShardingMetadata (class in orbax.checkpoint.metadata) ShardingTypes (class in orbax.checkpoint.metadata) should_save() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method) shutdown_dtensor() (in module orbax.export.dtensor_utils) signatures (orbax.export.utils.CallableSignatures property) SingleDeviceShardingMetadata (class in orbax.checkpoint.metadata) standard_name_format() (in module orbax.checkpoint.path.step) StandardCheckpointer (class in orbax.checkpoint) StandardCheckpointHandler (class in orbax.checkpoint) StandardRestore (in module orbax.checkpoint.args) StandardSave (in module orbax.checkpoint.args) step (orbax.checkpoint.path.step.Metadata attribute) step_from_checkpoint_name() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) step_metadata_of_checkpoint_path() (in module orbax.checkpoint.path.step) step_prefix_with_underscore() (in module orbax.checkpoint.path.step) StringHandler (class in orbax.checkpoint.type_handlers) StringMetadata (class in orbax.checkpoint.metadata) structure() (orbax.checkpoint.AbstractCheckpointer method) sync_global_processes() (in module orbax.checkpoint.multihost)

T

TensorSpecWithDefault (class in orbax.export.utils) tf_module (orbax.export.export_manager.ExportManager property) tmp_checkpoints() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.utils) to_flat_dict() (in module orbax.checkpoint.utils) to_jax_sharding() (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method) to_json() (orbax.checkpoint.metadata.TreeMetadata method) to_serialized_string() (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method)

Transform (class in orbax.checkpoint.transform_utils) TreeMetadata (class in orbax.checkpoint.metadata) TypeHandler (class in orbax.checkpoint.type_handlers) TypeHandlerRegistry (class in orbax.checkpoint.type_handlers) typestr() (orbax.checkpoint.type_handlers.ArrayHandler method), [1] (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.ScalarHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.type_handlers.TypeHandler method)

U

unlock_existing_checkpoints() (in module orbax.checkpoint.checkpoint_utils)

update_variables() (orbax.export.jax_module.JaxModule method)

V

validate() (orbax.export.validate.validation_manager.ValidationManager method) ValidationManager (class in orbax.export.validate.validation_manager)

ValidationReport (class in orbax.export.validate.validation_report) ValidationReportOption (class in orbax.export.validate.validation_report)

W

wait_for_new_checkpoint() (in module orbax.checkpoint.checkpoint_utils) wait_until_finished() (orbax.checkpoint.abstract_checkpoint_manager.AbstractCheckpointManager method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.checkpoint_manager.CheckpointManager method)

with_default_args() (in module orbax.export.utils)