Index _ | A | B | C | D | E | F | G | H | I | J | L | M | N | O | P | R | S | T | U | V | W | X _ __call__() (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving.ScopedStorageOptionsCreator method) (orbax.checkpoint.multihost.BarrierSyncFn method) __delattr__() (orbax.checkpoint.experimental.v1.training.RootMetadata method), [1] (orbax.checkpoint.metadata.StorageMetadata method) (orbax.checkpoint.options.FileOptions method) (orbax.checkpoint.path.step.Metadata method) __eq__() (orbax.checkpoint.args.CheckpointArgs method) (orbax.checkpoint.checkpoint_managers.preservation_policy.AnyPreservationPolicy method) (orbax.checkpoint.checkpoint_managers.preservation_policy.BestN method) (orbax.checkpoint.checkpoint_managers.preservation_policy.CustomSteps method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSeconds method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSteps method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNStepsClosest method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestDuration method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestN method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationContext method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreserveAll method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.AnySavePolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.ContinuousCheckpointingPolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.DecisionContext method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.FixedIntervalPolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.SpecificStepsPolicy method) (orbax.checkpoint.CheckpointManagerOptions method) (orbax.checkpoint.experimental.v1.options.ArrayOptions method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Loading method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Loading.LoadAndBroadcastOptions method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving.StorageOptions method) (orbax.checkpoint.experimental.v1.options.AsyncOptions method) (orbax.checkpoint.experimental.v1.options.CheckpointablesOptions method) (orbax.checkpoint.experimental.v1.options.DeletionOptions method) (orbax.checkpoint.experimental.v1.options.DeletionOptions.GcsDeletionOptions method) (orbax.checkpoint.experimental.v1.options.FileOptions method) (orbax.checkpoint.experimental.v1.options.MemoryOptions method) (orbax.checkpoint.experimental.v1.options.MultiprocessingOptions method) (orbax.checkpoint.experimental.v1.options.PathwaysOptions method) (orbax.checkpoint.experimental.v1.options.PyTreeOptions method) (orbax.checkpoint.experimental.v1.options.PyTreeOptions.Loading method) (orbax.checkpoint.experimental.v1.options.PyTreeOptions.Saving method) (orbax.checkpoint.experimental.v1.options.SafetensorsOptions method) (orbax.checkpoint.experimental.v1.training.preservation_policies.PreservationContext method) (orbax.checkpoint.experimental.v1.training.RootMetadata method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.DecisionContext method) (orbax.checkpoint.handlers.ArrayRestoreArgs method) (orbax.checkpoint.handlers.ArraySaveArgs method) (orbax.checkpoint.handlers.JaxRandomKeyRestoreArgs method) (orbax.checkpoint.handlers.JaxRandomKeySaveArgs method) (orbax.checkpoint.handlers.JsonRestoreArgs method) (orbax.checkpoint.handlers.JsonSaveArgs method) (orbax.checkpoint.handlers.NumpyRandomKeyRestoreArgs method) (orbax.checkpoint.handlers.NumpyRandomKeySaveArgs method) (orbax.checkpoint.handlers.ProtoRestoreArgs method) (orbax.checkpoint.handlers.ProtoSaveArgs method) (orbax.checkpoint.handlers.PyTreeRestoreArgs method) (orbax.checkpoint.handlers.PyTreeSaveArgs method) (orbax.checkpoint.handlers.StandardRestoreArgs method) (orbax.checkpoint.handlers.StandardSaveArgs method) (orbax.checkpoint.logging.CloudLoggerOptions method) (orbax.checkpoint.metadata.ArrayMetadata method) (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.ScalarMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method) (orbax.checkpoint.metadata.StorageMetadata method) (orbax.checkpoint.metadata.StringMetadata method) (orbax.checkpoint.options.AsyncOptions method) (orbax.checkpoint.options.FileOptions method) (orbax.checkpoint.options.MultiprocessingOptions method) (orbax.checkpoint.path.step.Metadata method) (orbax.checkpoint.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) __get_pydantic_core_schema__() (orbax.checkpoint.experimental.v1.path.Path class method) __hash__ (orbax.checkpoint.args.CheckpointArgs attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.AnyPreservationPolicy attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.BestN attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.CustomSteps attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSeconds attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSteps attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNStepsClosest attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestDuration attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestN attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationContext attribute) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreserveAll attribute) (orbax.checkpoint.checkpoint_managers.save_decision_policy.AnySavePolicy attribute) (orbax.checkpoint.checkpoint_managers.save_decision_policy.ContinuousCheckpointingPolicy attribute) (orbax.checkpoint.checkpoint_managers.save_decision_policy.DecisionContext attribute) (orbax.checkpoint.checkpoint_managers.save_decision_policy.FixedIntervalPolicy attribute) (orbax.checkpoint.checkpoint_managers.save_decision_policy.SpecificStepsPolicy attribute) (orbax.checkpoint.CheckpointManagerOptions attribute) (orbax.checkpoint.experimental.v1.options.ArrayOptions attribute) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Loading attribute) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Loading.LoadAndBroadcastOptions attribute) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving.StorageOptions attribute) (orbax.checkpoint.experimental.v1.options.AsyncOptions attribute) (orbax.checkpoint.experimental.v1.options.CheckpointablesOptions attribute) (orbax.checkpoint.experimental.v1.options.DeletionOptions attribute) (orbax.checkpoint.experimental.v1.options.DeletionOptions.GcsDeletionOptions attribute) (orbax.checkpoint.experimental.v1.options.FileOptions attribute) (orbax.checkpoint.experimental.v1.options.MemoryOptions attribute) (orbax.checkpoint.experimental.v1.options.MultiprocessingOptions attribute) (orbax.checkpoint.experimental.v1.options.PathwaysOptions attribute) (orbax.checkpoint.experimental.v1.options.PyTreeOptions attribute) (orbax.checkpoint.experimental.v1.options.PyTreeOptions.Loading attribute) (orbax.checkpoint.experimental.v1.options.PyTreeOptions.Saving attribute) (orbax.checkpoint.experimental.v1.options.SafetensorsOptions attribute) (orbax.checkpoint.experimental.v1.training.preservation_policies.PreservationContext attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.DecisionContext attribute) (orbax.checkpoint.handlers.ArrayRestoreArgs attribute) (orbax.checkpoint.handlers.ArraySaveArgs attribute) (orbax.checkpoint.handlers.JaxRandomKeyRestoreArgs attribute) (orbax.checkpoint.handlers.JaxRandomKeySaveArgs attribute) (orbax.checkpoint.handlers.JsonRestoreArgs attribute) (orbax.checkpoint.handlers.JsonSaveArgs attribute) (orbax.checkpoint.handlers.NumpyRandomKeyRestoreArgs attribute) (orbax.checkpoint.handlers.NumpyRandomKeySaveArgs attribute) (orbax.checkpoint.handlers.ProtoRestoreArgs attribute) (orbax.checkpoint.handlers.ProtoSaveArgs attribute) (orbax.checkpoint.handlers.PyTreeRestoreArgs attribute) (orbax.checkpoint.handlers.PyTreeSaveArgs attribute) (orbax.checkpoint.handlers.StandardRestoreArgs attribute) (orbax.checkpoint.handlers.StandardSaveArgs attribute) (orbax.checkpoint.logging.CloudLoggerOptions attribute) (orbax.checkpoint.metadata.ArrayMetadata attribute) (orbax.checkpoint.metadata.NamedShardingMetadata attribute) (orbax.checkpoint.metadata.ScalarMetadata attribute) (orbax.checkpoint.metadata.ShardingMetadata attribute) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata attribute) (orbax.checkpoint.metadata.StringMetadata attribute) (orbax.checkpoint.options.AsyncOptions attribute) (orbax.checkpoint.options.MultiprocessingOptions attribute) (orbax.checkpoint.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.experimental.v1.training.RootMetadata method), [1] (orbax.checkpoint.metadata.StorageMetadata method) (orbax.checkpoint.options.FileOptions method) (orbax.checkpoint.path.step.Metadata method) __init__() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.args.CheckpointArgs method) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.checkpoint_managers.preservation_policy.AnyPreservationPolicy method) (orbax.checkpoint.checkpoint_managers.preservation_policy.BestN method) (orbax.checkpoint.checkpoint_managers.preservation_policy.CustomSteps method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSeconds method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSteps method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNStepsClosest method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestDuration method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestN method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationContext method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationPolicy method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreserveAll method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.AnySavePolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.ContinuousCheckpointingPolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.DecisionContext method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.FixedIntervalPolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.SaveDecisionPolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.SpecificStepsPolicy method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointManager method) (orbax.checkpoint.CheckpointManagerOptions method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.experimental.v1.CheckpointMetadata method) (orbax.checkpoint.experimental.v1.Context method) (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler method) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler method) (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler method) (orbax.checkpoint.experimental.v1.options.ArrayOptions method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Loading method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Loading.LoadAndBroadcastOptions method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving.ScopedStorageOptionsCreator method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving.StorageOptions method) (orbax.checkpoint.experimental.v1.options.AsyncOptions method) (orbax.checkpoint.experimental.v1.options.CheckpointablesOptions method) (orbax.checkpoint.experimental.v1.options.DeletionOptions method) (orbax.checkpoint.experimental.v1.options.DeletionOptions.GcsDeletionOptions method) (orbax.checkpoint.experimental.v1.options.FileOptions method) (orbax.checkpoint.experimental.v1.options.MemoryOptions method) (orbax.checkpoint.experimental.v1.options.MultiprocessingOptions method) (orbax.checkpoint.experimental.v1.options.PathwaysOptions method) (orbax.checkpoint.experimental.v1.options.PyTreeOptions method) (orbax.checkpoint.experimental.v1.options.PyTreeOptions.Loading method) (orbax.checkpoint.experimental.v1.options.PyTreeOptions.Saving method) (orbax.checkpoint.experimental.v1.options.SafetensorsOptions method) (orbax.checkpoint.experimental.v1.path.PathAwaitingCreation method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] (orbax.checkpoint.experimental.v1.training.CheckpointMetadata method), [1] (orbax.checkpoint.experimental.v1.training.preservation_policies.PreservationContext method) (orbax.checkpoint.experimental.v1.training.RootMetadata method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.DecisionContext method) (orbax.checkpoint.handlers.ArrayRestoreArgs method) (orbax.checkpoint.handlers.ArraySaveArgs method) (orbax.checkpoint.handlers.JaxRandomKeyRestoreArgs method) (orbax.checkpoint.handlers.JaxRandomKeySaveArgs method) (orbax.checkpoint.handlers.JsonRestoreArgs method) (orbax.checkpoint.handlers.JsonSaveArgs method) (orbax.checkpoint.handlers.NumpyRandomKeyRestoreArgs method) (orbax.checkpoint.handlers.NumpyRandomKeySaveArgs method) (orbax.checkpoint.handlers.ProtoRestoreArgs method) (orbax.checkpoint.handlers.ProtoSaveArgs method) (orbax.checkpoint.handlers.PyTreeRestoreArgs method) (orbax.checkpoint.handlers.PyTreeSaveArgs method) (orbax.checkpoint.handlers.StandardRestoreArgs method) (orbax.checkpoint.handlers.StandardSaveArgs method) (orbax.checkpoint.JsonCheckpointHandler method) (orbax.checkpoint.logging.CloudLogger method) (orbax.checkpoint.logging.CloudLoggerOptions method) (orbax.checkpoint.logging.CompositeLogger method) (orbax.checkpoint.metadata.ArrayMetadata method) (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.ScalarMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method) (orbax.checkpoint.metadata.StorageMetadata method) (orbax.checkpoint.metadata.StringMetadata method) (orbax.checkpoint.multihost.BarrierSyncFn method) (orbax.checkpoint.options.AsyncOptions method) (orbax.checkpoint.options.FileOptions method) (orbax.checkpoint.options.MultiprocessingOptions 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 method) (orbax.checkpoint.type_handlers.ArrayHandler method) (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.SingleReplicaArrayHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.v1.path.step.NameFormat 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) __new__() (orbax.checkpoint.AbstractCheckpointer static method) (orbax.checkpoint.experimental.v1.path.Path static method) __post_init__() (orbax.checkpoint.checkpoint_managers.preservation_policy.CustomSteps method) (orbax.export.serving_config.ServingConfig method) (orbax.export.validate.validation_report.ValidationReportOption method) __setattr__() (orbax.checkpoint.experimental.v1.training.RootMetadata method), [1] (orbax.checkpoint.metadata.StorageMetadata method) (orbax.checkpoint.options.FileOptions method) (orbax.checkpoint.path.step.Metadata method) __subclasshook__() (orbax.checkpoint.AbstractCheckpointManager class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.AnyPreservationPolicy class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.BestN class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.CustomSteps class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSeconds class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSteps class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNStepsClosest class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestDuration class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestN class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationPolicy class method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreserveAll class method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.AnySavePolicy class method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.ContinuousCheckpointingPolicy class method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.FixedIntervalPolicy class method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.InitialSavePolicy class method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.PreemptionCheckpointingPolicy class method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.SaveDecisionPolicy class method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.SpecificStepsPolicy class method) (orbax.checkpoint.CheckpointManager class method) (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler class method) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler class method) (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler class method) (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving.ScopedStorageOptionsCreator class method) (orbax.checkpoint.experimental.v1.path.PathAwaitingCreation class method) (orbax.checkpoint.multihost.BarrierSyncFn class method) (orbax.checkpoint.path.step.NameFormat class method) (orbax.checkpoint.v1.path.step.NameFormat class method) A absolute() (orbax.checkpoint.experimental.v1.path.Path method) AbstractArray (class in orbax.checkpoint.experimental.v1.arrays) AbstractCheckpointer (class in orbax.checkpoint) AbstractCheckpointManager (class in orbax.checkpoint) AbstractLeaf (in module orbax.checkpoint.experimental.v1.tree) AbstractLogger (class in orbax.checkpoint.logging) AbstractShardedArray (class in orbax.checkpoint.experimental.v1.arrays) active_processes (orbax.checkpoint.experimental.v1.options.MultiprocessingOptions attribute) add() (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandlerRegistry method) all_steps() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) any_checkpoint_step() (in module orbax.checkpoint.path.step) AnyPreservationPolicy (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) AnySavePolicy (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) apply_fn_map (orbax.export.jax_module.JaxModule property) apply_transformations() (in module orbax.checkpoint) array (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) array_metadata_store (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) array_metadata_validator (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler attribute), [1] ArrayCheckpointHandler (class in orbax.checkpoint) ArrayHandler (class in orbax.checkpoint.type_handlers) ArrayLike (in module orbax.checkpoint.experimental.v1.arrays) ArrayMetadata (class in orbax.checkpoint.metadata) ArrayOptions (class in orbax.checkpoint.experimental.v1.options) ArrayOptions.Loading (class in orbax.checkpoint.experimental.v1.options) ArrayOptions.Loading.LoadAndBroadcastOptions (class in orbax.checkpoint.experimental.v1.options) ArrayOptions.Saving (class in orbax.checkpoint.experimental.v1.options) ArrayOptions.Saving.ScopedStorageOptionsCreator (class in orbax.checkpoint.experimental.v1.options) ArrayOptions.Saving.StorageOptions (class in orbax.checkpoint.experimental.v1.options) ArrayRestore (in module orbax.checkpoint.args) ArrayRestoreArgs (class in orbax.checkpoint.handlers) (class in orbax.checkpoint.type_handlers) ArraySave (in module orbax.checkpoint.args) ArraySaveArgs (class in orbax.checkpoint.handlers) async_save() (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointHandler 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) AsyncCheckpointer (class in orbax.checkpoint) AsyncCheckpointHandler (class in orbax.checkpoint) asynchronous (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) AsyncOptions (class in orbax.checkpoint.experimental.v1.options) (class in orbax.checkpoint.options) AsyncResponse (class in orbax.checkpoint.v1) AtomicRenameTemporaryPath (class in orbax.checkpoint.path.atomicity) await_creation() (orbax.checkpoint.experimental.v1.path.PathAwaitingCreation method) B barrier_sync_key_prefix (orbax.checkpoint.experimental.v1.options.MultiprocessingOptions attribute) BarrierSyncFn (class in orbax.checkpoint.multihost) best_step() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) BestN (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) bind() (orbax.export.serving_config.ServingConfig method) broadcast_one_to_all() (in module orbax.checkpoint.multihost) build_name() (orbax.checkpoint.path.step.NameFormat method) (orbax.checkpoint.v1.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.AbstractCheckpointManager method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.CheckpointManager method) check_input() (orbax.export.validate.validation_manager.ValidationManager class method) check_output() (orbax.export.validate.validation_manager.ValidationManager class method) checkpoint_layout (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) checkpoint_steps() (in module orbax.checkpoint.path.step) checkpoint_steps_paths() (in module orbax.checkpoint.path.step) CheckpointableHandler (class in orbax.checkpoint.experimental.v1.handlers), [1] (class in orbax.checkpoint.v1) CheckpointableHandlerRegistry (class in orbax.checkpoint.experimental.v1.handlers) checkpointables (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) checkpointables_metadata() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] CheckpointablesOptions (class in orbax.checkpoint.experimental.v1.options) CheckpointArgs (class in orbax.checkpoint.args) Checkpointer (class in orbax.checkpoint) (class in orbax.checkpoint.experimental.v1.training), [1] CheckpointHandler (class in orbax.checkpoint) checkpointing_impl (orbax.checkpoint.experimental.v1.options.PathwaysOptions attribute) CheckpointLayout (class in orbax.checkpoint.experimental.v1.options) CheckpointManager (class in orbax.checkpoint) CheckpointManagerOptions (class in orbax.checkpoint) CheckpointMetadata (class in orbax.checkpoint.experimental.v1) (class in orbax.checkpoint.experimental.v1.training), [1] (class in orbax.checkpoint.v1) checkpoints (orbax.checkpoint.experimental.v1.training.Checkpointer property), [1] checkpoints_iterator() (in module orbax.checkpoint.checkpoint_utils) close() (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CheckpointManager method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) CloudLogger (class in orbax.checkpoint.logging) CloudLoggerOptions (class in orbax.checkpoint.logging) commit_timestamp (orbax.checkpoint.experimental.v1.CheckpointMetadata property) (orbax.checkpoint.path.step.Metadata property) commit_timestamp_nsecs (orbax.checkpoint.experimental.v1.CheckpointMetadata attribute) (orbax.checkpoint.path.step.Metadata property) (orbax.checkpoint.v1.CheckpointMetadata attribute) CommitFileTemporaryPath (class in orbax.checkpoint.path.atomicity) Composite (in module orbax.checkpoint.args) composite_name_format() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.v1.path.step) CompositeArgs (class in orbax.checkpoint.handlers) CompositeCheckpointHandler (class in orbax.checkpoint) CompositeLogger (class in orbax.checkpoint.logging) CompositeResults (in module orbax.checkpoint.handlers) construct_restore_args() (in module orbax.checkpoint.checkpoint_utils) Context (class in orbax.checkpoint.experimental.v1) (class in orbax.checkpoint.v1) context (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler attribute), [1] ContinuousCheckpointingPolicy (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) copy() (orbax.checkpoint.experimental.v1.path.Path method) create() (orbax.checkpoint.path.atomicity.AtomicRenameTemporaryPath method) (orbax.checkpoint.path.atomicity.CommitFileTemporaryPath method) (orbax.checkpoint.path.atomicity_types.TemporaryPath method) create_directories_asynchronously (orbax.checkpoint.experimental.v1.options.AsyncOptions attribute) create_type_handler_registry() (in module orbax.checkpoint.type_handlers) custom_metadata (orbax.checkpoint.experimental.v1.CheckpointMetadata attribute) (orbax.checkpoint.experimental.v1.training.RootMetadata attribute), [1] (orbax.checkpoint.handlers.PyTreeSaveArgs attribute) (orbax.checkpoint.handlers.StandardSaveArgs attribute) (orbax.checkpoint.v1.CheckpointMetadata attribute) CustomSteps (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) D data_processors (orbax.export.serving_config.ServingConfig attribute) DecisionContext (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) default_val (orbax.export.utils.TensorSpecWithDefault attribute) delete() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) deletion (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) DeletionOptions (class in orbax.checkpoint.experimental.v1.options) DeletionOptions.GcsDeletionOptions (class in orbax.checkpoint.experimental.v1.options) deserialize() (orbax.checkpoint.type_handlers.ArrayHandler method) (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.ScalarHandler method) (orbax.checkpoint.type_handlers.SingleReplicaArrayHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.type_handlers.TypeHandler method) deserialize_tree() (in module orbax.checkpoint.tree) directory (orbax.checkpoint.AbstractCheckpointManager property) (orbax.checkpoint.CheckpointManager property) (orbax.checkpoint.experimental.v1.training.Checkpointer property), [1] (orbax.checkpoint.experimental.v1.training.RootMetadata attribute), [1] dtensor_initialized() (in module orbax.export.dtensor_utils) E enable_pinned_host_transfer (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) enable_post_merge_validation (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) enable_replica_parallel_separate_folder (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) enable_write_sharding_file (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) EveryNSeconds (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) EveryNSteps (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) EveryNStepsClosest (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) exists() (orbax.checkpoint.experimental.v1.path.Path method) expanduser() (orbax.checkpoint.experimental.v1.path.Path method) export_module() (orbax.export.jax_module.JaxModule method) export_version (orbax.export.jax_module.JaxModule property) ExportManager (class in orbax.export.export_manager) extra_trackable_resources (orbax.export.serving_config.ServingConfig attribute) F file (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) filename (orbax.checkpoint.experimental.v1.handlers.JsonHandler attribute) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler attribute), [1] FileOptions (class in orbax.checkpoint.experimental.v1.options) (class in orbax.checkpoint.options) finalize() (in module orbax.checkpoint.experimental.v1.partial) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.path.atomicity.AtomicRenameTemporaryPath method) (orbax.checkpoint.path.atomicity.CommitFileTemporaryPath method) (orbax.checkpoint.path.atomicity_types.TemporaryPath method) (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) (orbax.checkpoint.type_handlers.TypeHandler method) find_all() (orbax.checkpoint.path.step.NameFormat method) (orbax.checkpoint.v1.path.step.NameFormat method) find_step() (orbax.checkpoint.path.step.NameFormat method) (orbax.checkpoint.v1.path.step.NameFormat method) find_step_path() (in module orbax.checkpoint.path.step) FixedIntervalPolicy (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) floating_atol (orbax.export.validate.validation_report.ValidationReportOption attribute) floating_rtol (orbax.export.validate.validation_report.ValidationReportOption attribute) format() (orbax.checkpoint.experimental.v1.path.Path method) from_deserialized_dict() (orbax.checkpoint.metadata.NamedShardingMetadata class method) (orbax.checkpoint.metadata.ShardingMetadata class method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata class method) from_final() (orbax.checkpoint.path.atomicity.AtomicRenameTemporaryPath class method) (orbax.checkpoint.path.atomicity.CommitFileTemporaryPath class method) (orbax.checkpoint.path.atomicity_types.TemporaryPath class method) from_flat_dict() (in module orbax.checkpoint.tree) from_jax_sharding() (orbax.checkpoint.metadata.NamedShardingMetadata class method) (orbax.checkpoint.metadata.ShardingMetadata class method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata class method) from_saved_model() (orbax.export.utils.CallableSignatures class method) from_temporary() (orbax.checkpoint.path.atomicity.AtomicRenameTemporaryPath class method) (orbax.checkpoint.path.atomicity.CommitFileTemporaryPath class method) (orbax.checkpoint.path.atomicity_types.TemporaryPath class method) G gcs_deletion_options (orbax.checkpoint.experimental.v1.options.DeletionOptions attribute) get() (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandlerRegistry method) (orbax.checkpoint.path.atomicity_types.TemporaryPath method) get_all_entries() (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandlerRegistry method) get_barrier_sync_fn() (in module orbax.checkpoint.multihost) 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_final() (orbax.checkpoint.path.atomicity.AtomicRenameTemporaryPath method) (orbax.checkpoint.path.atomicity.CommitFileTemporaryPath method) (orbax.checkpoint.path.atomicity_types.TemporaryPath method) 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_param_names() (in module orbax.checkpoint.tree) 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) get_ts_context() (in module orbax.checkpoint.type_handlers) get_type_handler() (in module orbax.checkpoint.type_handlers) glob() (orbax.checkpoint.experimental.v1.path.Path method) global_registry() (in module orbax.checkpoint.experimental.v1.handlers), [1] H has() (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandlerRegistry method) has_registered_args() (in module orbax.checkpoint.args) has_type_handler() (in module orbax.checkpoint.type_handlers) I ignore_load_sharding (orbax.checkpoint.experimental.v1.options.SafetensorsOptions attribute) init_timestamp (orbax.checkpoint.experimental.v1.CheckpointMetadata property) init_timestamp_nsecs (orbax.checkpoint.experimental.v1.CheckpointMetadata attribute) (orbax.checkpoint.path.step.Metadata property) (orbax.checkpoint.v1.CheckpointMetadata attribute) initialize_dtensor() (in module orbax.export.dtensor_utils) InitialSavePolicy (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) input_polymorphic_shape_map (orbax.export.jax_module.JaxModule property) input_signature (orbax.export.serving_config.ServingConfig attribute) interval (orbax.checkpoint.checkpoint_managers.save_decision_policy.FixedIntervalPolicy attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.FixedIntervalPolicy attribute) InvalidLayoutError (class in orbax.checkpoint.experimental.v1.errors) is_abstract_handleable() (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler method) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler method) (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler method) is_async_checkpointer() (in module orbax.checkpoint.checkpoint_manager) is_checkpoint_finalized() (in module orbax.checkpoint.path.step) is_dir() (orbax.checkpoint.experimental.v1.path.Path method) is_file() (orbax.checkpoint.experimental.v1.path.Path method) is_handleable() (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler method) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler method) (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler method) is_ocdbt_checkpoint() (in module orbax.checkpoint.type_handlers) is_orbax_checkpoint() (in module orbax.checkpoint.v1) is_pathways_backend() (in module orbax.checkpoint.experimental.v1.multihost) is_primary (orbax.export.utils.TensorSpecWithDefault attribute) is_primary_host() (in module orbax.checkpoint.experimental.v1.multihost) (in module orbax.checkpoint.multihost) is_prioritized_key_fn (orbax.checkpoint.experimental.v1.options.MemoryOptions attribute) is_saving_in_progress (orbax.checkpoint.checkpoint_managers.save_decision_policy.DecisionContext attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.DecisionContext attribute) is_saving_in_progress() (orbax.checkpoint.CheckpointManager method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] is_tmp_checkpoint() (in module orbax.checkpoint.path.step) item (orbax.checkpoint.handlers.ArrayRestoreArgs attribute) (orbax.checkpoint.handlers.ArraySaveArgs attribute) (orbax.checkpoint.handlers.JaxRandomKeySaveArgs attribute) (orbax.checkpoint.handlers.JsonRestoreArgs attribute) (orbax.checkpoint.handlers.JsonSaveArgs attribute) (orbax.checkpoint.handlers.NumpyRandomKeySaveArgs attribute) (orbax.checkpoint.handlers.ProtoRestoreArgs attribute) (orbax.checkpoint.handlers.ProtoSaveArgs attribute) (orbax.checkpoint.handlers.PyTreeSaveArgs attribute) (orbax.checkpoint.handlers.StandardSaveArgs attribute) item_metadata() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) iterdir() (orbax.checkpoint.experimental.v1.path.Path method) J jax2tf_kwargs_map (orbax.export.jax_module.JaxModule property) 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) JaxRandomKeyRestoreArgs (class in orbax.checkpoint.handlers) JaxRandomKeySave (in module orbax.checkpoint.args) JaxRandomKeySaveArgs (class in orbax.checkpoint.handlers) JsonCheckpointHandler (class in orbax.checkpoint) JsonHandler (class in orbax.checkpoint.experimental.v1.handlers) JsonRestore (in module orbax.checkpoint.args) JsonRestoreArgs (class in orbax.checkpoint.handlers) JsonSave (in module orbax.checkpoint.args) JsonSaveArgs (class in orbax.checkpoint.handlers) JsonType (in module orbax.checkpoint.experimental.v1.tree) L latency (orbax.export.validate.validation_report.ValidationReport attribute) latest (orbax.checkpoint.experimental.v1.training.Checkpointer property), [1] latest_step() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) latest_step_metadata() (in module orbax.checkpoint.path.step) LatestDuration (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) LatestN (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) Leaf (in module orbax.checkpoint.experimental.v1.tree) listdir() (orbax.checkpoint.experimental.v1.path.Path method) load() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler method) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler method) (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] (orbax.export.export_manager.ExportManager method) load_async() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] load_checkpointables() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] load_checkpointables_async() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] loading (orbax.checkpoint.experimental.v1.options.ArrayOptions attribute) (orbax.checkpoint.experimental.v1.options.PyTreeOptions attribute) local_registry() (in module orbax.checkpoint.experimental.v1.handlers), [1] log_entry() (orbax.checkpoint.logging.AbstractLogger method) (orbax.checkpoint.logging.CloudLogger method) (orbax.checkpoint.logging.CompositeLogger method) (orbax.checkpoint.logging.StandardLogger method) M make_auto_batching_function() (in module orbax.export.utils) max_non_floating_mismatch_ratio (orbax.export.validate.validation_report.ValidationReportOption attribute) max_replicas_for_replica_parallel (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) maybe_enable_dtensor_export_on() (in module orbax.export.dtensor_utils) memory (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) memory_size() (orbax.checkpoint.type_handlers.ArrayHandler method) (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.ScalarHandler method) (orbax.checkpoint.type_handlers.SingleReplicaArrayHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.type_handlers.TypeHandler method) MemoryOptions (class in orbax.checkpoint.experimental.v1.options) merge() (in module orbax.checkpoint.experimental.v1.tree) merge_ocdbt_per_process_files() (in module orbax.checkpoint.type_handlers) merge_trees() (in module orbax.checkpoint) Metadata (class in orbax.checkpoint.path.step) metadata (orbax.checkpoint.experimental.v1.CheckpointMetadata attribute) (orbax.checkpoint.v1.CheckpointMetadata attribute) (orbax.export.validate.validation_report.ValidationReport attribute) metadata() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CheckpointManager method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler method) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler method) (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointHandler method) (orbax.checkpoint.type_handlers.ArrayHandler method) (orbax.checkpoint.type_handlers.NumpyHandler method) (orbax.checkpoint.type_handlers.ScalarHandler method) (orbax.checkpoint.type_handlers.StringHandler method) (orbax.checkpoint.type_handlers.TypeHandler method) metadata_from_temporary_paths() (orbax.checkpoint.CompositeCheckpointHandler method) method_key (orbax.export.serving_config.ServingConfig attribute) methods (orbax.export.jax_module.JaxModule property) metrics() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) min_slice_bytes_for_replica_parallel (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) minimum_interval_secs (orbax.checkpoint.checkpoint_managers.save_decision_policy.ContinuousCheckpointingPolicy attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.ContinuousCheckpointingPolicy attribute) mkdir() (orbax.checkpoint.experimental.v1.path.Path method) model_param_names (orbax.export.jax_module.JaxModule property) model_params (orbax.export.jax_module.JaxModule property) module orbax.checkpoint orbax.checkpoint.abstract_checkpoint_manager orbax.checkpoint.args orbax.checkpoint.arrays orbax.checkpoint.checkpoint_manager orbax.checkpoint.checkpoint_managers.preservation_policy orbax.checkpoint.checkpoint_managers.save_decision_policy orbax.checkpoint.checkpoint_utils orbax.checkpoint.experimental.v1, [1], [2], [3] orbax.checkpoint.experimental.v1.arrays orbax.checkpoint.experimental.v1.errors orbax.checkpoint.experimental.v1.handlers, [1], [2], [3], [4] orbax.checkpoint.experimental.v1.multihost orbax.checkpoint.experimental.v1.options orbax.checkpoint.experimental.v1.partial orbax.checkpoint.experimental.v1.path orbax.checkpoint.experimental.v1.training, [1], [2] orbax.checkpoint.experimental.v1.training.errors orbax.checkpoint.experimental.v1.training.preservation_policies orbax.checkpoint.experimental.v1.training.save_decision_policies orbax.checkpoint.experimental.v1.tree orbax.checkpoint.handlers orbax.checkpoint.handlers.handler_registration orbax.checkpoint.logging orbax.checkpoint.metadata orbax.checkpoint.msgpack_utils orbax.checkpoint.multihost orbax.checkpoint.options orbax.checkpoint.path.atomicity orbax.checkpoint.path.atomicity_types orbax.checkpoint.path.deleter orbax.checkpoint.path.step orbax.checkpoint.transform_utils orbax.checkpoint.tree orbax.checkpoint.type_handlers orbax.export.dtensor_utils orbax.export.export_manager 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) multiprocessing (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) multiprocessing_options (orbax.checkpoint.checkpoint_managers.save_decision_policy.DecisionContext attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.DecisionContext attribute) MultiprocessingOptions (class in orbax.checkpoint.experimental.v1.options) (class in orbax.checkpoint.options) N NamedShardingMetadata (class in orbax.checkpoint.metadata) NameFormat (class in orbax.checkpoint.path.step) (class in orbax.checkpoint.v1.path.step) NumpyHandler (class in orbax.checkpoint.type_handlers) NumpyRandomKeyCheckpointHandler (class in orbax.checkpoint) NumpyRandomKeyRestore (in module orbax.checkpoint.args) NumpyRandomKeyRestoreArgs (class in orbax.checkpoint.handlers) NumpyRandomKeySave (in module orbax.checkpoint.args) NumpyRandomKeySaveArgs (class in orbax.checkpoint.handlers) O obm_export_options (orbax.export.serving_config.ServingConfig attribute) obm_module_to_jax_exported_map() (orbax.export.jax_module.JaxModule method) ocdbt_target_data_file_size (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) (orbax.checkpoint.handlers.PyTreeSaveArgs attribute) on_commit_callback() (in module orbax.checkpoint.path.atomicity) open() (orbax.checkpoint.experimental.v1.path.Path method) orbax.checkpoint module orbax.checkpoint.abstract_checkpoint_manager module orbax.checkpoint.args module orbax.checkpoint.arrays module orbax.checkpoint.checkpoint_manager module orbax.checkpoint.checkpoint_managers.preservation_policy module orbax.checkpoint.checkpoint_managers.save_decision_policy module orbax.checkpoint.checkpoint_utils module orbax.checkpoint.experimental.v1 module, [1], [2], [3] orbax.checkpoint.experimental.v1.arrays module orbax.checkpoint.experimental.v1.errors module orbax.checkpoint.experimental.v1.handlers module, [1], [2], [3], [4] orbax.checkpoint.experimental.v1.multihost module orbax.checkpoint.experimental.v1.options module orbax.checkpoint.experimental.v1.partial module orbax.checkpoint.experimental.v1.path module orbax.checkpoint.experimental.v1.training module, [1], [2] orbax.checkpoint.experimental.v1.training.errors module orbax.checkpoint.experimental.v1.training.preservation_policies module orbax.checkpoint.experimental.v1.training.save_decision_policies module orbax.checkpoint.experimental.v1.tree module orbax.checkpoint.handlers module orbax.checkpoint.handlers.handler_registration module orbax.checkpoint.logging module orbax.checkpoint.metadata module orbax.checkpoint.msgpack_utils module orbax.checkpoint.multihost module orbax.checkpoint.options module orbax.checkpoint.path.atomicity module orbax.checkpoint.path.atomicity_types module orbax.checkpoint.path.deleter module orbax.checkpoint.path.step module orbax.checkpoint.transform_utils module orbax.checkpoint.tree module orbax.checkpoint.type_handlers module orbax.export.dtensor_utils module orbax.export.export_manager 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 orbax_module() (orbax.export.jax_module.JaxModule method) output_report_path (orbax.export.validate.validation_report.ValidationReportOption attribute) outputs (orbax.export.validate.validation_report.ValidationReport attribute) P Path (class in orbax.checkpoint.experimental.v1.path) path (orbax.checkpoint.experimental.v1.CheckpointMetadata property) (orbax.checkpoint.experimental.v1.training.errors.StepAlreadyExistsError attribute) (orbax.checkpoint.experimental.v1.training.errors.StepNotFoundError attribute) (orbax.checkpoint.path.step.Metadata attribute) path_class (orbax.checkpoint.experimental.v1.options.FileOptions attribute), [1] path_permission_mode (orbax.checkpoint.experimental.v1.options.FileOptions attribute) (orbax.checkpoint.options.FileOptions attribute) PathAwaitingCreation (class in orbax.checkpoint.experimental.v1.path) PathLike (in module orbax.checkpoint.experimental.v1.path) pathways (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) PathwaysOptions (class in orbax.checkpoint.experimental.v1.options) PLACEHOLDER (in module orbax.checkpoint.v1) policies (orbax.checkpoint.checkpoint_managers.save_decision_policy.AnySavePolicy attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.AnySavePolicy attribute) post_finalization_callback (orbax.checkpoint.experimental.v1.options.AsyncOptions attribute) postprocessors (orbax.export.serving_config.ServingConfig attribute) PreemptionCheckpointingPolicy (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) preprocess_output_passthrough_enabled (orbax.export.serving_config.ServingConfig attribute) preprocessors (orbax.export.serving_config.ServingConfig attribute) PreservationContext (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) PreservationPolicy (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) PreserveAll (class in orbax.checkpoint.checkpoint_managers.preservation_policy) (class in orbax.checkpoint.experimental.v1.training.preservation_policies) primary_host (orbax.checkpoint.experimental.v1.options.MultiprocessingOptions attribute) print_debug_info (orbax.export.validate.validation_report.ValidationReportOption attribute) process_count() (in module orbax.checkpoint.experimental.v1.multihost) process_index() (in module orbax.checkpoint.experimental.v1.multihost) (in module orbax.checkpoint.multihost) ProtoCheckpointHandler (class in orbax.checkpoint) ProtoHandler (class in orbax.checkpoint.experimental.v1.handlers), [1] ProtoRestore (in module orbax.checkpoint.args) ProtoRestoreArgs (class in orbax.checkpoint.handlers) ProtoSave (in module orbax.checkpoint.args) ProtoSaveArgs (class in orbax.checkpoint.handlers) PyTree (in module orbax.checkpoint.experimental.v1.tree) pytree (orbax.checkpoint.experimental.v1.Context attribute) (orbax.checkpoint.v1.Context attribute) PyTreeCheckpointer (class in orbax.checkpoint) PyTreeCheckpointHandler (class in orbax.checkpoint) PyTreeHandler (class in orbax.checkpoint.experimental.v1.handlers), [1] PyTreeKey (in module orbax.checkpoint.experimental.v1.tree) PyTreeKeyPath (in module orbax.checkpoint.experimental.v1.tree) PyTreeMetadata (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) PyTreeOf (in module orbax.checkpoint.experimental.v1.tree) PyTreeOptions (class in orbax.checkpoint.experimental.v1.options) PyTreeOptions.Loading (class in orbax.checkpoint.experimental.v1.options) PyTreeOptions.Saving (class in orbax.checkpoint.experimental.v1.options) PyTreeRestore (in module orbax.checkpoint.args) PyTreeRestoreArgs (class in orbax.checkpoint.handlers) PyTreeSave (in module orbax.checkpoint.args) PyTreeSaveArgs (class in orbax.checkpoint.handlers) R reached_preemption (orbax.checkpoint.checkpoint_managers.save_decision_policy.DecisionContext attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.DecisionContext attribute) reached_preemption() (in module orbax.checkpoint.multihost) (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) read_bytes() (orbax.checkpoint.experimental.v1.path.Path method) read_concurrent_bytes (orbax.checkpoint.experimental.v1.options.MemoryOptions attribute) read_text() (orbax.checkpoint.experimental.v1.path.Path method) register_handler() (in module orbax.checkpoint.experimental.v1.handlers), [1] 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) registry (orbax.checkpoint.experimental.v1.options.CheckpointablesOptions attribute) reload() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.CheckpointManager method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] rename() (orbax.checkpoint.experimental.v1.path.Path method) replace() (orbax.checkpoint.experimental.v1.path.Path method) resolve() (orbax.checkpoint.experimental.v1.path.Path method) restore() (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CheckpointManager 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.StandardCheckpointer method) (orbax.checkpoint.StandardCheckpointHandler method) restore_args (orbax.checkpoint.handlers.ArrayRestoreArgs attribute) (orbax.checkpoint.handlers.JaxRandomKeyRestoreArgs attribute) restore_type (orbax.checkpoint.type_handlers.ArrayRestoreArgs attribute) RestoreArgs (class in orbax.checkpoint.type_handlers) rglob() (orbax.checkpoint.experimental.v1.path.Path method) rmdir() (orbax.checkpoint.experimental.v1.path.Path method) rmtree() (orbax.checkpoint.experimental.v1.path.Path method) RootMetadata (class in orbax.checkpoint.experimental.v1.training), [1] S SafetensorsOptions (class in orbax.checkpoint.experimental.v1.options) save() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.experimental.v1.partial) (in module orbax.checkpoint.v1) (orbax.checkpoint.AbstractCheckpointer method) (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.ArrayCheckpointHandler method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.Checkpointer method) (orbax.checkpoint.CheckpointHandler method) (orbax.checkpoint.CheckpointManager method) (orbax.checkpoint.CompositeCheckpointHandler method) (orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler method) (orbax.checkpoint.experimental.v1.handlers.ProtoHandler method) (orbax.checkpoint.experimental.v1.handlers.PyTreeHandler method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] (orbax.checkpoint.JaxRandomKeyCheckpointHandler method) (orbax.checkpoint.JsonCheckpointHandler method) (orbax.checkpoint.NumpyRandomKeyCheckpointHandler method) (orbax.checkpoint.ProtoCheckpointHandler method) (orbax.checkpoint.PyTreeCheckpointHandler method) (orbax.checkpoint.StandardCheckpointer method) (orbax.checkpoint.StandardCheckpointHandler method) (orbax.export.export_manager.ExportManager method) save_args (orbax.checkpoint.handlers.ArraySaveArgs attribute) (orbax.checkpoint.handlers.JaxRandomKeySaveArgs attribute) (orbax.checkpoint.handlers.PyTreeSaveArgs attribute) (orbax.checkpoint.handlers.StandardSaveArgs attribute) save_async() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.experimental.v1.partial) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] save_checkpointables() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] save_checkpointables_async() (in module orbax.checkpoint.experimental.v1) (in module orbax.checkpoint.v1) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] save_shlo_to_file (orbax.export.jax_module.JaxModule property) SaveArgs (class in orbax.checkpoint.type_handlers) SaveDecisionPolicy (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) saving (orbax.checkpoint.experimental.v1.options.ArrayOptions attribute) (orbax.checkpoint.experimental.v1.options.PyTreeOptions attribute) ScalarHandler (class in orbax.checkpoint.type_handlers) ScalarMetadata (class in orbax.checkpoint.metadata) scoped_storage_options_creator (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) serialize() (orbax.checkpoint.type_handlers.ArrayHandler method) (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.tree) serving_signatures (orbax.export.export_manager.ExportManager property) ServingConfig (class in orbax.export.serving_config) ShardingMetadata (class in orbax.checkpoint.metadata) should_preserve() (orbax.checkpoint.checkpoint_managers.preservation_policy.AnyPreservationPolicy method) (orbax.checkpoint.checkpoint_managers.preservation_policy.BestN method) (orbax.checkpoint.checkpoint_managers.preservation_policy.CustomSteps method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSeconds method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSteps method) (orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNStepsClosest method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestDuration method) (orbax.checkpoint.checkpoint_managers.preservation_policy.LatestN method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationPolicy method) (orbax.checkpoint.checkpoint_managers.preservation_policy.PreserveAll method) (orbax.checkpoint.experimental.v1.training.preservation_policies.AnyPreservationPolicy method) (orbax.checkpoint.experimental.v1.training.preservation_policies.BestN method) (orbax.checkpoint.experimental.v1.training.preservation_policies.CustomSteps method) (orbax.checkpoint.experimental.v1.training.preservation_policies.EveryNSeconds method) (orbax.checkpoint.experimental.v1.training.preservation_policies.EveryNSteps method) (orbax.checkpoint.experimental.v1.training.preservation_policies.EveryNStepsClosest method) (orbax.checkpoint.experimental.v1.training.preservation_policies.LatestDuration method) (orbax.checkpoint.experimental.v1.training.preservation_policies.LatestN method) (orbax.checkpoint.experimental.v1.training.preservation_policies.PreservationPolicy method) (orbax.checkpoint.experimental.v1.training.preservation_policies.PreserveAll method) should_save() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.AnySavePolicy method), [1] (orbax.checkpoint.checkpoint_managers.save_decision_policy.ContinuousCheckpointingPolicy method), [1] (orbax.checkpoint.checkpoint_managers.save_decision_policy.FixedIntervalPolicy method), [1] (orbax.checkpoint.checkpoint_managers.save_decision_policy.InitialSavePolicy method), [1] (orbax.checkpoint.checkpoint_managers.save_decision_policy.PreemptionCheckpointingPolicy method), [1] (orbax.checkpoint.checkpoint_managers.save_decision_policy.SaveDecisionPolicy method) (orbax.checkpoint.checkpoint_managers.save_decision_policy.SpecificStepsPolicy method), [1] (orbax.checkpoint.CheckpointManager method) (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.AnySavePolicy method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.ContinuousCheckpointingPolicy method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.FixedIntervalPolicy method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.InitialSavePolicy method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.PreemptionCheckpointingPolicy method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.SaveDecisionPolicy method), [1] (orbax.checkpoint.experimental.v1.training.save_decision_policies.SpecificStepsPolicy method), [1] shutdown_dtensor() (in module orbax.export.dtensor_utils) signature_key (orbax.export.serving_config.ServingConfig attribute) signatures (orbax.export.utils.CallableSignatures property) SingleDeviceShardingMetadata (class in orbax.checkpoint.metadata) SingleReplicaArrayHandler (class in orbax.checkpoint.type_handlers) SpecificStepsPolicy (class in orbax.checkpoint.checkpoint_managers.save_decision_policy) (class in orbax.checkpoint.experimental.v1.training.save_decision_policies) standard_name_format() (in module orbax.checkpoint.path.step) (in module orbax.checkpoint.v1.path.step) StandardCheckpointer (class in orbax.checkpoint) StandardCheckpointHandler (class in orbax.checkpoint) StandardLogger (class in orbax.checkpoint.logging) StandardRestore (in module orbax.checkpoint.args) StandardRestoreArgs (class in orbax.checkpoint.handlers) StandardSave (in module orbax.checkpoint.args) StandardSaveArgs (class in orbax.checkpoint.handlers) stat() (orbax.checkpoint.experimental.v1.path.Path method) StatefulCheckpointable (class in orbax.checkpoint.experimental.v1.handlers) (class in orbax.checkpoint.v1) status (orbax.export.validate.validation_report.ValidationReport attribute) step (orbax.checkpoint.experimental.v1.training.errors.StepAlreadyExistsError attribute) (orbax.checkpoint.experimental.v1.training.errors.StepNotFoundError attribute) (orbax.checkpoint.path.step.Metadata attribute) step_from_checkpoint_name() (in module orbax.checkpoint.path.step) step_metadata_of_checkpoint_path() (in module orbax.checkpoint.path.step) step_prefix_with_underscore() (in module orbax.checkpoint.path.step) StepAlreadyExistsError (class in orbax.checkpoint.experimental.v1.training.errors) StepNotFoundError (class in orbax.checkpoint.experimental.v1.training.errors) steps (orbax.checkpoint.checkpoint_managers.save_decision_policy.SpecificStepsPolicy attribute) (orbax.checkpoint.experimental.v1.training.save_decision_policies.SpecificStepsPolicy attribute) storage_options (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) StorageMetadata (class in orbax.checkpoint.metadata) 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.experimental.v1.multihost) (in module orbax.checkpoint.multihost) T temporary_path_class (orbax.checkpoint.experimental.v1.options.FileOptions attribute) TemporaryPath (class in orbax.checkpoint.path.atomicity_types) tensor_spec (orbax.export.utils.TensorSpecWithDefault attribute) TensorSpecWithDefault (class in orbax.export.utils) tf_module (orbax.export.export_manager.ExportManager property) tf_postprocessor (orbax.export.serving_config.ServingConfig attribute) tf_preprocessor (orbax.export.serving_config.ServingConfig attribute) timeout_secs (orbax.checkpoint.experimental.v1.options.AsyncOptions attribute) tmp_checkpoints() (in module orbax.checkpoint.path.step) to_flat_dict() (in module orbax.checkpoint.tree) to_jax_sharding() (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method) to_serialized_string() (orbax.checkpoint.metadata.NamedShardingMetadata method) (orbax.checkpoint.metadata.ShardingMetadata method) (orbax.checkpoint.metadata.SingleDeviceShardingMetadata method) to_shape_dtype_struct() (in module orbax.checkpoint.arrays) (in module orbax.checkpoint.experimental.v1.arrays) (in module orbax.checkpoint.tree) todelete_full_path (orbax.checkpoint.experimental.v1.options.DeletionOptions.GcsDeletionOptions attribute) touch() (orbax.checkpoint.experimental.v1.path.Path method) transfer_concurrent_bytes (orbax.checkpoint.experimental.v1.options.MemoryOptions attribute) Transform (class in orbax.checkpoint) TypeHandler (class in orbax.checkpoint.type_handlers) TypeHandlerRegistry (class in orbax.checkpoint.type_handlers) typestr() (orbax.checkpoint.CheckpointHandler class method) (orbax.checkpoint.type_handlers.ArrayHandler method) (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 unlink() (orbax.checkpoint.experimental.v1.path.Path method) update_variables() (orbax.export.jax_module.JaxModule method) use_compression (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) use_ocdbt (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) use_replica_parallel (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) use_zarr3 (orbax.checkpoint.experimental.v1.options.ArrayOptions.Saving attribute) V v0() (orbax.checkpoint.experimental.v1.options.FileOptions method) validate() (orbax.checkpoint.path.atomicity.AtomicRenameTemporaryPath class method) (orbax.checkpoint.path.atomicity.CommitFileTemporaryPath class method) (orbax.checkpoint.path.atomicity_types.TemporaryPath class method) (orbax.export.validate.validation_manager.ValidationManager method) validate_final() (orbax.checkpoint.path.atomicity.AtomicRenameTemporaryPath class method) (orbax.checkpoint.path.atomicity.CommitFileTemporaryPath class method) (orbax.checkpoint.path.atomicity_types.TemporaryPath class 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() (orbax.checkpoint.experimental.v1.training.Checkpointer method), [1] wait_for_new_checkpoint() (in module orbax.checkpoint.checkpoint_utils) wait_until_finished() (orbax.checkpoint.AbstractCheckpointManager method) (orbax.checkpoint.AsyncCheckpointer method) (orbax.checkpoint.CheckpointManager method) with_default_args() (in module orbax.export.utils) with_gradient (orbax.export.jax_module.JaxModule property) write_bytes() (orbax.checkpoint.experimental.v1.path.Path method) write_concurrent_bytes (orbax.checkpoint.experimental.v1.options.MemoryOptions attribute) write_text() (orbax.checkpoint.experimental.v1.path.Path method) X xprof_url (orbax.export.validate.validation_report.ValidationReport attribute)