Orbax

https://orbax.readthedocs.io/en/latest/
Admin
Orbax is an umbrella namespace providing common training utilities for JAX users, including flexible...
Tokens:201,362
Snippets:1,271
Trust Score:8.3
Update:3 months ago
Show doc for...
Tokens:
Raw
### Defining Dummy PyTrees and Metadata

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html

Setup of dummy PyTrees and metadata structures for checkpointing examples.

```python
# Dummy PyTrees for simplicity.

# In reality, this would be a tree of np.ndarray or jax.Array.
pytree = {'a': 0}
# In reality, this would be a tree of jax.ShapeDtypeStruct (metadata
# for restoration).
abstract_pytree = {'a': 0}

extra_metadata = {'version': 1.0}
```

--------------------------------

### Install Pre-commit Hooks

Source: https://orbax.readthedocs.io/en/latest/contributors.html

Install the pre-commit framework to ensure code style compliance before committing changes.

```bash
pip install pre-commit
pre-commit install
```

--------------------------------

### Initialize Checkpoint Environment

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_api_overview.html

Setup imports and create a temporary directory for checkpointing.

```python
import orbax.checkpoint as ocp
from orbax.checkpoint.checkpoint_managers import preservation_policy as preservation_policy_lib
from orbax.checkpoint.checkpoint_managers import save_decision_policy as save_decision_policy_lib
import jax
import numpy as np
from jax import numpy as jnp

path = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')

state = {'layer0': {'bias': np.ones((4,)), 'weight': jnp.arange(16)}}
abstract_state = jax.tree.map(ocp.tree.to_shape_dtype_struct, state)
metadata = {'version': 1.0}
extra_metadata = {'version': 1.0, 'step': 0}
dataset = {'my_data': 2}
```

--------------------------------

### Install orbax-checkpoint from PyPI

Source: https://orbax.readthedocs.io/en/latest/installation.html

Standard installation command for the latest stable release of the checkpointing package.

```bash
pip install orbax-checkpoint
```

--------------------------------

### Install orbax-export from PyPI

Source: https://orbax.readthedocs.io/en/latest/installation.html

Standard installation command for the latest stable release of the exporting package.

```bash
pip install orbax-export
```

--------------------------------

### Install orbax-checkpoint from GitHub

Source: https://orbax.readthedocs.io/en/latest/installation.html

Use this command to install the most recent version directly from the source repository.

```bash
pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'
```

--------------------------------

### Log Output Examples

Source: https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html

Standard error logs generated during initialization.

```text
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1784214086.525412    7910 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
I0000 00:00:1784214086.537787    7910 cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
```

```text
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1784214088.543110    7910 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
I0000 00:00:1784214088.543643    7910 cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
```

--------------------------------

### Install orbax-export from GitHub

Source: https://orbax.readthedocs.io/en/latest/installation.html

Use this command to install the exporting package directly from the source repository.

```bash
pip install 'git+https://github.com/google/orbax/#subdirectory=export'
```

--------------------------------

### Setup Training State

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/async_checkpointing.html

Initializes the necessary imports and a sample training state dictionary for checkpointing.

```python
import numpy as np
import orbax.checkpoint as ocp
from etils import epath

train_state = {
    'layer0': {
        'kernel': np.ones((8, 8), dtype=np.float32),
        'bias': np.ones((8,), dtype=np.float32),
    }
}
```

--------------------------------

### Enable preprocess output passthrough

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/export/serving_config.html

Example implementation of pre/post-processors when preprocess_output_passthrough_enabled is set to True.

```python
def tf_preprocessor(x):
  return {'pre_out_to_jax': x}, {'pre_out_to_post': x}

def jax_func(inputs: Mapping[str, tf.Tensor]) -> Mapping[str, tf.Tensor]:
  return {'jax_out_to_post': inputs['pre_out_to_jax']}

def tf_postprocessor(
    inputs: Mapping[str, tf.Tensor],
    inputs_extra: Mapping[str, tf.Tensor],
) -> Mapping[str, tf.Tensor]:
  return {
      'post_out_from_jax': inputs['jax_out_to_post'],
      'post_out_from_pre': inputs_extra['pre_out_to_post'],
  }
preprocess_output_passthrough_enabled: bool = False
```

--------------------------------

### Initialize CheckpointManager

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_101.html

Setup the directory and initial state variables for checkpoint management.

```python
path = ocp.test_utils.erase_and_create_empty('/tmp/checkpoint_manager')
state = {
    'a': np.arange(8),
    'b': np.arange(16),
}
extra_params = [42, 43]
```

--------------------------------

### Initialize PyTreeCheckpointHandler

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/checkpointing_pytrees.html

Setup the AsyncCheckpointer with a PyTreeCheckpointHandler for advanced serialization.

```python
path = ocp.test_utils.erase_and_create_empty('/tmp/advanced/')
# Make sure to use async for improved performance!
ckptr = ocp.AsyncCheckpointer(ocp.PyTreeCheckpointHandler())
```

--------------------------------

### Example Output for Subsequent Restore

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_api_overview.html

Output showing the restored state and None for the omitted extra_metadata.

```text
{'layer0': {'bias': array([11., 11., 11., 11.]), 'weight': Array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25],      dtype=int32)}}
None
```

--------------------------------

### Example Parameter Output

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/debug_guide.html

Sample output representing the structure of parameters in a checkpoint.

```text
a.x
a.y
b
c
```

--------------------------------

### Install orbax-checkpoint from a Pull Request

Source: https://orbax.readthedocs.io/en/latest/installation.html

Install and test changes from a specific GitHub pull request by replacing the placeholder with the PR ID.

```bash
pip install 'git+https://github.com/google/orbax.git@refs/pull/<PR_NUMBER>/head#subdirectory=checkpoint'
```

--------------------------------

### Checkpoint metadata output

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_101.html

Example output showing detailed metadata for a saved checkpoint.

```text
StepMetadata(item_handlers='orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler', item_metadata=TreeMetadata(
  custom_metadata=None
  tree={'a': ArrayMetadata :  name=a,  directory=/tmp/my-checkpoints/checkpoint_name,  shape=(8,),  sharding=None,  dtype=int64,  storage=StorageMetadata(chunk_shape=(8,), write_shape=None),, 'b': {'c': ScalarMetadata(name='b.c', directory=PosixGPath('/tmp/my-checkpoints/checkpoint_name'), shape=(), sharding=None, dtype=dtype('int64'), storage=None), 'd': ArrayMetadata :  name=b.d,  directory=/tmp/my-checkpoints/checkpoint_name,  shape=(16,),  sharding=None,  dtype=int64,  storage=StorageMetadata(chunk_shape=(16,), write_shape=None),}}
  use_zarr3=False
), metrics={}, performance_metrics=SaveStepStatistics(step=None, event_type='save', directory=None, reached_preemption=False, preemption_received_at=None, synchronous=False, wait_for_prev_start_time=None, wait_for_prev_duration_secs=None, time_between_consecutive_saves_sec=None, checkpointer_blocking_start_time=None, checkpointer_blocking_duration_secs=None, get_old_steps_start_time=None, get_old_steps_duration_secs=None, checkpoint_manager_blocking_start_time=None, checkpoint_manager_blocking_duration_secs=None), init_timestamp_nsecs=1784213869896591548, commit_timestamp_nsecs=1784213869931872389, custom_metadata={})
```

--------------------------------

### Checkpoint directory contents output

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_101.html

Example output showing the sub-directories created for the composite checkpoint.

```python
[PosixGPath('/tmp/my-checkpoints/composite_checkpoint/_CHECKPOINT_METADATA'),
 PosixGPath('/tmp/my-checkpoints/composite_checkpoint/metadata'),
 PosixGPath('/tmp/my-checkpoints/composite_checkpoint/state')]
```

--------------------------------

### Initialize LatestDuration Policy

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/checkpoint_managers/preservation_policy.html

Example of initializing the LatestDuration policy to retain checkpoints from the last 24 hours.

```python
import datetime
LatestDuration(datetime.timedelta(hours=24))
```

--------------------------------

### Example Warning Output for Missing Handler

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_api_overview.html

Warning messages generated when the handler is not configured.

```text
WARNING:absl:Item "extra_metadata" was found in the checkpoint, but could not be restored. Please provide a `CheckpointHandlerRegistry`, or call `restore` with an appropriate `CheckpointArgs` subclass.
```

```text
WARNING:absl:Item "state" was found in the checkpoint, but could not be restored. Please provide a `CheckpointHandlerRegistry`, or call `restore` with an appropriate `CheckpointArgs` subclass.
```

```text
WARNING:absl:Item "extra_metadata" was found in the checkpoint, but could not be restored. Please provide a `CheckpointHandlerRegistry`, or call `restore` with an appropriate `CheckpointArgs` subclass.
```

```text
WARNING:absl:Item "state" was found in the checkpoint, but could not be restored. Please provide a `CheckpointHandlerRegistry`, or call `restore` with an appropriate `CheckpointArgs` subclass.
```

--------------------------------

### Checkpoint Directory Structure Output

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/checkpoint_format.html

Example output showing the directory tree of a saved state, including manifest files and data subdirectories.

```text
Printing directory tree: /tmp/checkpoint/0/state/
..manifest.ocdbt
..d/
....e2cf3b93d0dcb4aeb7052bcc392a9ac7
..ocdbt.process_0/
....manifest.ocdbt
....d/
......8de146f57b26a2471369cd6f22a84c77
......7cca1002a70192f1c2ea029d76dc45d0
......badda37ad130ce65bb5946e7065350e3
......d831ad77a255444550d559dffac0103d
.._strings.json
..array_metadatas/
....process_0
.._sharding
.._METADATA
```

--------------------------------

### Setup JAX Model and Pre/Post-processing

Source: https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html

Define the JAX model function, parameters, and optional TensorFlow pre/post-processing functions required for export.

```python
# Import Orbax Export classes.
from orbax.export import ExportManager
from orbax.export import JaxModule
from orbax.export import ServingConfig
import numpy as np
import jax
import jax.numpy as jnp
import tensorflow as tf

# Prepare the parameters and model function to export.
example1_params = {  'a': np.array(5.0), 'b': np.array(1.1), 'c': np.array(0.55)} # A pytree of the JAX model parameters.

# model f(x) = a * sin(x) + b * x + c, here (a, b, c) are model parameters
def example1_model_fn(params, inputs):  # The JAX model function to export.
  a, b, c = params['a'], params['b'], params['c']
  return a * jnp.sin(inputs) + b * inputs + c

def example1_preprocess(inputs):  # Optional: preprocessor in TF.
  norm_inputs = tf.nest.map_structure(lambda x: x/tf.math.reduce_max(x), inputs)
  return norm_inputs

def example1_postprocess(model_fn_outputs):  # Optional: post-processor in TF.
  return {'outputs': model_fn_outputs}
```

--------------------------------

### Setup imports for Orbax checkpointing

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/checkpoint_format.html

Initial imports required to interact with Orbax checkpointing and file system utilities.

```python
from etils import epath
import jax
import numpy as np
import orbax.checkpoint as ocp
```

--------------------------------

### CompositeCheckpointHandler Usage Examples

Source: https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.handlers.html

Demonstrates basic usage without a registry, handling errors with mismatched arguments, and configuring a custom handler registry.

```python
# The simplest use-case, with no handler registry provided on construction.
checkpointer = ocp.Checkpointer(
    ocp.CompositeCheckpointHandler()
)

checkpointer.save(directory,
    ocp.args.Composite(
        # The handler will be determined from the global registry using the
        # provided args.
        state=ocp.args.StandardSave(pytree),
    )
)

restored: ocp.args.Composite = checkpoint_handler.restore(directory,
    ocp.args.Composite(
        state=ocp.args.StandardRestore(abstract_pytree),
    )
)
restored.state ...

# The "state" item's args was determined on the first call to `save`.
# Trying to save "state" with a different set of args will raise an error.
checkpointer.save(directory,
  ocp.args.Composite(
    # Will raise `ValueError: Item "state" and args "JsonSave [...]" does
    # not match with any registered handler! ...`.
    state=ocp.args.JsonSave(pytree),
    )
)

# You can also register specific items and handlers using a handler
# registry.
handler_registry = (
    ocp.handler_registration.DefaultCheckpointHandlerRegistry()
)
# Some subclass of `CheckpointHandler` that handles the "state" item.
handler = state_handler()
# Multiple save args and handlers can be registered for the same item.
handler_registry.add(
    'state,
    StateSaveArgs,
    handler,
)
handler_registry.add(
    'state',
    StateRestoreArgs,
    handler,
)

# The `CompositeCheckpointHandler` will use the handlers registered in the
# registry when it encounters the item name `state`.
checkpointer_with_registry = ocp.Checkpointer(
    ocp.CompositeCheckpointHandler(handler_registry=handler_registry)
)

checkpointer_with_registry.save(directory,
    ocp.args.Composite(
        # The handler knows that the "state" item should use the
        # `state_handler` specified in the registry. When an item has been
        # added to the handler registry, only save args that are registered
        # for that item are allowed.
        state=StateSaveArgs(pytree),
        # You can also provide other arguments that are not registered in
        # the handler registry as long as a handler has been globally
        # registered for the args.
        other_param=ocp.args.StandardSave(other_pytree),
    )
)

restored: ocp.args.Composite = checkpointer_with_registry.restore(directory,
    ocp.args.Composite(
        state=StateRestoreArgs(abstract_pytree),
        other_param=ocp.args.StandardRestore(abstract_other_pytree),
    )
)
restored.state ...
restored.other_param...

# Skip restoring `other_params` (you can save a subset of items too, in a
# similar fashion).
restored_state_only: ocp.args.Composite = (
  checkpointer_with_registry.restore(directory,
    ocp.args.Composite(
        state=StateRestoreArgs(abstract_pytree),
    )
  )
)
restored.state ...
restored.other_param ... # None.
```

--------------------------------

### Create a ServingConfig with Processors

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/export/serving_config.html

Demonstrates initializing a ServingConfig with custom TensorFlow-based pre- and post-processors.

```python
import tensorflow as tf
from orbax.export import tf_data_processor
from orbax.export.serving_config import ServingConfig

@tf.function(input_signature=[tf.TensorSpec(
    shape=(None, 32), dtype=tf.float32
)])
def preprocessor(inputs):
  return {'normalized': inputs / 255.0}

def postprocessor(outputs):
  return {'probabilities': tf.nn.softmax(outputs)}

config = ServingConfig(
    signature_key='serving_default',
    preprocessors=[tf_data_processor.TfDataProcessor(preprocessor)],
    postprocessors=[tf_data_processor.TfDataProcessor(postprocessor)]
)
```

--------------------------------

### ServingConfig.__init__

Source: https://orbax.readthedocs.io/en/latest/api_reference/export.serving_config.html

Initializes a new ServingConfig instance with signature keys, input signatures, and various pre/post-processors.

```APIDOC
## ServingConfig.__init__

### Description
Initializes the ServingConfig object with configuration parameters for model serving.

### Parameters
- **signature_key** (str) - Required - The key for the serving signature.
- **input_signature** (Any) - Optional - The input signature definition.
- **tf_preprocessor** (Callable) - Optional - TensorFlow preprocessor function.
- **tf_postprocessor** (Callable) - Optional - TensorFlow postprocessor function.
- **preprocessors** (tuple) - Optional - List of preprocessors.
- **postprocessors** (tuple) - Optional - List of postprocessors.
- **data_processors** (tuple) - Optional - List of data processors.
- **extra_trackable_resources** (Any) - Optional - Additional trackable resources.
- **method_key** (str) - Optional - The method key.
- **obm_export_options** (Any) - Optional - Options for OBM export.
- **preprocess_output_passthrough_enabled** (bool) - Optional - Whether to enable preprocess output passthrough.
```

--------------------------------

### Execute the benchmark

Source: https://orbax.readthedocs.io/en/latest/checkpointing_benchmark_guide.html

Run the benchmark suite from the repository root by specifying the configuration file and output directory.

```bash
python checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py \
  --config_file=sweep.yaml --output_directory=/tmp/bench/
```

--------------------------------

### Initialize Checkpointer

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/checkpointing_pytrees.html

Set up the AsyncCheckpointer with a StandardCheckpointHandler.

```python
path = ocp.test_utils.erase_and_create_empty('/tmp/basic/')
# Make sure to use async for improved performance!
ckptr = ocp.AsyncCheckpointer(ocp.StandardCheckpointHandler())
```

--------------------------------

### Save PyTree Example

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.html

Example usage of the Checkpointer with PyTreeCheckpointHandler to save a nested dictionary structure.

```python
      ckptr = ocp.Checkpointer(ocp.PyTreeCheckpointHandler())
      item = {
          'layer0': {
              'w': np.ndarray(...),
              'b': np.ndarray(...),
          },
          'layer1': {
              'w': np.ndarray(...),
              'b': np.ndarray(...),
          },
      }
      # Note: save_args may be None if no customization is desired for saved
      # parameters.
      # Otherwise, settings can be used to customize save behavior, e.g.
      # casting.
      save_args = jax.tree.map(lambda x: ocp.SaveArgs(dtype=np.int32), item)
      # Eventually calls through to `async_save`.
      ckptr.save(path, args=ocp.PyTreeSave(item, save_args))
```

--------------------------------

### Create a sample checkpoint

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/debug_guide.html

Initializes a sample state dictionary and saves it to a temporary path using StandardCheckpointer.

```python
state = {
    'a': {
        'x': np.arange(2 ** 24),
        'y': np.arange(1024),
    },
    'b': np.ones(8),
    'c': 42,
}

default_param_name = 'a.x'
default_path = epath.Path('/tmp/checkpoint')
if default_path.exists():
  default_path.rmtree()
with ocp.StandardCheckpointer() as ckptr:
  ckptr.save(default_path, state)
```

--------------------------------

### Initialize StandardCheckpointHandler

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/handlers/standard_checkpoint_handler.html

Demonstrates how to instantiate a Checkpointer using the StandardCheckpointHandler.

```python
ckptr = Checkpointer(StandardCheckpointHandler())
# OR
ckptr = StandardCheckpointer()
```

--------------------------------

### Metadata File Content Example

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/checkpoint_format.html

Example JSON content of a _METADATA file representing a flattened PyTree structure.

```json
{'tree_metadata': {"('a',)": {'key_metadata': [{'key': 'a', 'key_type': 2}],
   'value_metadata': {'value_type': 'jax.Array',
    'skip_deserialize': False,
    'write_shape': [16]}},
  "('b',)": {'key_metadata': [{'key': 'b', 'key_type': 2}],
   'value_metadata': {'value_type': 'jax.Array',
    'skip_deserialize': False,
    'write_shape': [16]}},
  "('c',)": {'key_metadata': [{'key': 'c', 'key_type': 2}],
   'value_metadata': {'value_type': 'np.ndarray', 'skip_deserialize': False}},
  "('d',)": {'key_metadata': [{'key': 'd', 'key_type': 2}],
   'value_metadata': {'value_type': 'scalar', 'skip_deserialize': False}},
  "('e',)": {'key_metadata': [{'key': 'e', 'key_type': 2}],
   'value_metadata': {'value_type': 'string', 'skip_deserialize': False}}},
 'use_ocdbt': True,
 'use_zarr3': False,
 'store_array_data_equal_to_fill_value': True,
 'custom_metadata': None}
```

--------------------------------

### Initialize step naming and metadata

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_manager.html

Configures step naming formats, loads existing checkpoint information, and initializes root metadata.

```python
    self._step_name_format = (
        self._options.step_name_format
        or step_lib.standard_name_format(
            step_prefix=self._options.step_prefix,
            step_format_fixed_length=self._options.step_format_fixed_length,
            single_host_load_and_broadcast=(
                self._options.single_host_load_and_broadcast
            ),
        )
    )

    self._checkpoints = checkpoint_info.CheckpointInfos(
        self._load_checkpoint_infos(
            skip_metadata_read=self._options.lightweight_initialize
        )
    )

    self._metadata_dir = self.directory / METADATA_ITEM_NAME
    if self._options.read_only and not self._metadata_dir.exists():
      custom_metadata = {} if metadata is None else dict(metadata)
    else:
      custom_metadata = None
    self._root_metadata = threading_lib.Ref(
        RootMetadata(
            custom_metadata=custom_metadata,
        )
    )

    self._maybe_save_root_metadata(metadata)
```

--------------------------------

### get_param_names

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/tree/utils.html

Gets parameter names for PyTree elements.

```APIDOC
## get_param_names(item, *, include_empty_nodes=True)

### Description
Generates parameter names for each element within a PyTree.

### Parameters
- **item** (PyTree) - Required - The PyTree to process.
- **include_empty_nodes** (bool) - Optional - Whether to include empty nodes in the result.

### Returns
- **PyTree** - A PyTree containing the parameter names.
```

--------------------------------

### Initialize StandardCheckpointer

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/checkpointers/standard_checkpointer.html

Demonstrates the shorthand initialization of StandardCheckpointer compared to the standard AsyncCheckpointer pattern.

```python
# Instead of:
with AsyncCheckpointer(StandardCheckpointHandler()) as ckptr:
  ...
# We can use:
with StandardCheckpointer() as ckptr:
  ...
```

--------------------------------

### Model output verification

Source: https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html

Example output from the loaded SavedModel.

```text
loaded model output:  tf.Tensor(
[ 0.99577844 -1.2953885  -0.9818985  -0.26964682 -0.8935896   3.5859964
 -1.3445401  -1.5546808   0.16264628  0.88514185 -0.7923399   1.800378
  0.63770825 -1.5368788  -1.7442648  -1.5777712 ], shape=(16,), dtype=float32)
```

--------------------------------

### Example Output for Item Metadata

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_api_overview.html

The structure of the returned TreeMetadata object.

```text
TreeMetadata(
  custom_metadata=None
  tree={'layer0': {'bias': ArrayMetadata :  name=layer0.bias,  directory=/tmp/checkpoint-manager-items-1/1/state,  shape=(4,),  sharding=None,  dtype=float64,  storage=StorageMetadata(chunk_shape=(4,), write_shape=None),, 'weight': ArrayMetadata :  name=layer0.weight,  directory=/tmp/checkpoint-manager-items-1/1/state,  shape=(16,),  sharding=SingleDeviceShardingMetadata(device_str=cpu:0),  dtype=int32,  storage=StorageMetadata(chunk_shape=(16,), write_shape=(16,)),}}
  use_zarr3=False
)
```

--------------------------------

### Restored metadata output

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_101.html

Example output structure of the restored metadata.

```python
{'version': 1.0, 'lang': 'en'}
```

--------------------------------

### Initialize Checkpoint Directory

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_api_overview.html

Sets up a temporary directory for the checkpoint manager.

```python
directory = ocp.test_utils.erase_and_create_empty('/tmp/checkpoint-manager-multiple/')
```

--------------------------------

### Initialize PyTreeCheckpointer

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/checkpointers/pytree_checkpointer.html

Demonstrates the shorthand usage of PyTreeCheckpointer compared to the standard Checkpointer initialization.

```python
ckptr = Checkpointer(PyTreeCheckpointHandler())
```

```python
ckptr = PyTreeCheckpointer()
```

--------------------------------

### Legacy Single-Item Checkpointing

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html

Example of the deprecated CheckpointManager usage pattern.

```python
options = ocp.CheckpointManagerOptions()
mngr = ocp.CheckpointManager(
  ocp.test_utils.erase_and_create_empty('/tmp/ckpt1/'),
  ocp.Checkpointer(ocp.PyTreeCheckpointHandler()),
  options=options,
)

restore_args = ocp.checkpoint_utils.construct_restore_args(abstract_pytree)
mngr.save(0, pytree)
mngr.wait_until_finished()

mngr.restore(
    0,
    items=abstract_pytree,
    restore_kwargs={'restore_args': restore_args}
)
```

--------------------------------

### Retrieve Checkpoint Metadata

Source: https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.handlers.html

Example of the metadata structure returned for a saved checkpoint.

```python
{
  'layer0': {
      'w': ArrayMetadata(dtype=jnp.float32, shape=(8, 8), shards=(1, 2)),
      'b': ArrayMetadata(dtype=jnp.float32, shape=(8,), shards=(1,)),
  },
  'step': ScalarMetadata(dtype=jnp.int64),
}
```

--------------------------------

### Using PathAwaitingCreation Protocol

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/path/types.html

Demonstrates logical access to a path and blocking until the path exists.

```python
path: :py:class:`.PathAwaitingCreation` = ...
# Logical accesses are OK.
print(path.path)
# Block until the path is known to exist.
path = await path.await_creation()
path.exists()  # True.
```

--------------------------------

### Define metadata dictionary

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_101.html

Example dictionary containing basic metadata for checkpointing.

```python
metadata = {
    'version': 1.0,
    'lang': 'en',
}
```

--------------------------------

### Using CompositeCheckpointHandler for Saving and Restoring

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.html

Demonstrates how to initialize a Checkpointer with CompositeCheckpointHandler and use Composite arguments to manage multiple data types simultaneously.

```python
import orbax.checkpoint as ocp

# A PyTree of jax.Arrays
my_state = ...
# A dictionary to be serialized as JSON
json_dict = ...

checkpointer = ocp.Checkpointer(ocp.CompositeCheckpointHandler())
checkpointer.save(
    path,
    args=ocp.args.Composite(
        state=ocp.args.StandardSave(my_state),
        metadata=ocp.args.JsonSave(json_dict)
    )
)

restored = checkpointer.restore(
    path,
    args=ocp.args.Composite(
        state=ocp.args.StandardRestore(),
        metadata=ocp.args.JsonRestore()
    )
)
my_state = restored.state
json_dict = restored.metadata
```

--------------------------------

### Example metadata structure

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.html

Representation of the metadata structure returned by the metadata method.

```python
      {
        'layer0': {
            'w': ArrayMetadata(dtype=jnp.float32, shape=(8, 8), shards=(1, 2)),
            'b': ArrayMetadata(dtype=jnp.float32, shape=(8,), shards=(1,)),
        },
        'step': ScalarMetadata(dtype=jnp.int64),
      }
```

--------------------------------

### Legacy Checkpointer Initialization

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_manager.html

Handles initialization for legacy-style checkpointers with specific constraints on primary host configuration.

```python
def _configure_checkpointer_legacy_init(
      self,
      checkpointers: Union[AbstractCheckpointer, CheckpointersDict],
      options: CheckpointManagerOptions,
  ) -> Checkpointer:
    """Initializes _CompositeCheckpointer with legacy style checkpointers."""
    if self._multiprocessing_options.primary_host != 0:
      raise ValueError(
          f'`primary_host`={self._multiprocessing_options.primary_host} is not'
          ' supported in legacy API.'
      )
```

--------------------------------

### Get Step Metadata

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_manager.html

Retrieves and validates metadata for a specific checkpoint step.

```python
def _get_step_metadata(self, step: int) -> StepMetadata:
  step_metadata = self._checkpointer.metadata(
      self._get_read_step_directory(step, self.directory),
  )

  self._default_item.set_if_none(
      _determine_default_item_mode_from_directory(
          self._get_read_step_directory(step, self.directory)
      )
  )
  step_metadata.item_metadata = self._maybe_get_default_item(
      step_metadata.item_metadata  # pyrefly: ignore[bad-argument-type]
  )
  if METRIC_ITEM_NAME in step_metadata.item_handlers:  # pyrefly: ignore[not-iterable]
    metrics = self._get_metrics(step)
    if metrics is not None:
      validated_metrics = step_metadata_serialization.deserialize(
          {}, metrics=dict(metrics)
      ).metrics
      step_metadata = dataclasses.replace(
          step_metadata,
          metrics=validated_metrics,
      )

  return step_metadata
```

--------------------------------

### Get Input Signature

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/export/serving_config.html

Retrieves the input signature from the configuration or the associated tf_preprocessor.

```python
  def get_input_signature(self, required=True) -> Any:
    """Gets the input signature from the explict one or tf_preprocessor."""
    input_signature = self.input_signature
    if input_signature is None:
      if (
          hasattr(self.tf_preprocessor, 'input_signature')
          and self.tf_preprocessor.input_signature is not None
      ):
        input_signature = self.tf_preprocessor.input_signature

    if required and input_signature is None:
      raise ValueError(
          f'Neithr the ServingConfig (key={self.signature_key}) nor its'
          ' `tf_preprocessor` sets an `input_signature`, please set'
          ' `input_signature` explicitly.',
      )
    return input_signature
```

--------------------------------

### Initialize and use CheckpointManager

Source: https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpoint_manager.html

Demonstrates saving and restoring checkpoints with multiple items or a single item using the CheckpointManager context manager.

```python
# Multiple items.
with CheckpointManager(
  'path/to/dir/',
  # Global metadata.
  metadata={'version': 1.1, 'lang': 'en'},
) as mngr:
  mngr.save(
      0,
      args=args.Composite(
        train_state=args.StandardSave(train_state),
        json_states=args.JsonSave(json_states),
      ),
      # Metadata varying by step.
      custom_metadata={'learning_rate': 0.001}
  )
  restored = mngr.restore(0)
  print(restored.train_state)
  print(restored.json_states)
  restored = mngr.restore(0, args=args.Composite(
      train_state=args.StandardRestore(abstract_train_state),
    )
  )
  print(restored.train_state)
  print(restored.json_states)  # Error, not restored
  global_metadata = mngr.metadata()
  step_metadata = mngr.metadata(0)  # custom_metadata in here

# Single, unnamed (default) item.
with CheckpointManager(
    'path/to/dir/',
    options = CheckpointManagerOptions(max_to_keep=5, ...),
  ) as mngr:
  mngr.save(0, args=StandardSave(train_state))
  train_state = mngr.restore(0)
  train_state = mngr.restore(0,
  args=StandardRestore(abstract_train_state))
```

--------------------------------

### Initialize ServingConfig func_map

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/export/serving_config.html

This snippet demonstrates how the inference function is mapped to signature keys.

```python
    func_map = {}
    infer_fn_with_processors = make_inference_fn(
        self.get_infer_step(infer_step_fns)
    )
    for key in self.get_signature_keys():
      func_map[key] = infer_fn_with_processors

    return func_map
```

--------------------------------

### Execute Benchmark Runner

Source: https://orbax.readthedocs.io/en/latest/checkpointing_benchmark_guide.html

Run the benchmark suite using the provided configuration file and output directory.

```bash
python checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py \
  --config_file=<config.yaml> \
  --output_directory=<dir> \
  [flags...]
```

--------------------------------

### Using Custom CheckpointArgs

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_args.html

Example of passing custom CheckpointArgs subclasses to the save method.

```python
ckptr.save(
    path,
    custom_state=MyCheckpointSave(item=..., options=...)
)

ckptr.save(
    path,
    custom_state=MyCheckpointRestore(options=...)
)
```

--------------------------------

### Visualize results in TensorBoard

Source: https://orbax.readthedocs.io/en/latest/checkpointing_benchmark_guide.html

Launch TensorBoard pointing to the output directory to view metrics, throughput, and memory usage.

```bash
tensorboard --logdir=/tmp/bench/tensorboard/
```

--------------------------------

### Get Metrics

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_manager.html

Restores metrics from the checkpoint directory using the metrics handler directly.

```python
def _get_metrics(self, step: int) -> PyTree | None:
  try:
    # Use handler directly, since this happens in a background thread and
    # barriers cannot be used. This usage pattern is not
    # recommended in other contexts.
    metrics = self._metrics_handler.restore(
        self._get_read_step_directory(step, self.directory) / METRIC_ITEM_NAME
    )
    return metrics
  except FileNotFoundError:
    logging.log_first_n(
        logging.INFO,
        'Missing metrics for step %d',
        1,  # log only once
        step,
    )
    return None
```

--------------------------------

### Checkpointer.__init__

Source: https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpointers.html

Initializes a new Checkpointer instance with a specific handler and optional configuration.

```APIDOC
## Checkpointer.__init__

### Description
Initializes the Checkpointer with a required CheckpointHandler and optional configuration for multiprocessing and file operations.

### Parameters
- **handler** (CheckpointHandler) - Required - The handler used for saving and restoring data.
- **multiprocessing_options** (MultiprocessingOptions) - Optional - Configuration for multi-process synchronization.
- **file_options** (FileOptions) - Optional - Configuration for file system permissions.
- **checkpoint_metadata_store** (Any) - Optional - Metadata store configuration.
- **temporary_path_class** (Any) - Optional - Class used for temporary path management.
```

--------------------------------

### Wrap JAX function and parameters into a JaxModule

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/export/jax_module.html

Demonstrates initializing a JaxModule with JAX parameters and an application function for export.

```python
import jax
import jax.numpy as jnp
from orbax.export import JaxModule

# Define parameters and a model application function
params = {'weights': jnp.ones((3, 3))}

def my_apply_fn(params, inputs):
  return jnp.dot(inputs, params['weights'])

# Create the exportable module (Defaults to TF_SAVEDMODEL format)
jax_module = JaxModule(
    params=params,
    apply_fn=my_apply_fn
)
```

--------------------------------

### Get metadata from checkpoint path

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/path/step.html

Resolves the metadata associated with a specific checkpoint path.

```python
def step_metadata_of_checkpoint_path(
    checkpoint_path: epath.PathLike, name_format: NameFormat[MetadataT]
) -> MetadataT:
  """Returns step.MetadataT of given `checkpoint_path`."""
  checkpoint_path = epath.Path(checkpoint_path)
  all_step_metadata = list(name_format.find_all(checkpoint_path.parent))
  for step_metadata in all_step_metadata:
    if step_metadata.path.name == checkpoint_path.name:
      return step_metadata
  raise ValueError(
      'Failed to resolve step metadata of checkpoint path with'
      f' NameFormat={name_format}, checkpoint path={checkpoint_path}, path'
      f' name({checkpoint_path.name}) did not match with available step names:'
      f' {[step_metadata.path.name for step_metadata in all_step_metadata]}.'
      ' Please check if the given path is really a checkpoint path.'
  )
```

--------------------------------

### Configure and export a JAX module

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/export/export_manager.html

Demonstrates initializing an ExportManager with a serving configuration and saving the model to a specified directory.

```python
import tensorflow as tf
from orbax.export import ExportManager
from orbax.export import serving_config

# Assume `my_jax_module` is a fully initialized jax_module.JaxModule
# Define how the model should handle incoming requests
my_config = serving_config.ServingConfig(
    signature_key="serving_default",
    input_signature=[tf.TensorSpec(shape=(None, 32), dtype=tf.float32)],
)

# Initialize the manager
export_mgr = ExportManager(
    module=my_jax_module,
    serving_configs=[my_config]
)

# Save the model to a directory
export_mgr.save("/path/to/my/saved_model")
```

--------------------------------

### Displaying Composite and TreeMetadata

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/orbax_checkpoint_api_overview.html

Example output showing the structure of Composite objects and associated TreeMetadata.

```text
Composite({'extra_metadata': {'version': 1.0, 'step': 4}, 'state': {'layer0': {'bias': array([11., 11., 11., 11.]), 'weight': Array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25],      dtype=int32)}}})

Composite({'extra_metadata': None, 'state': TreeMetadata(
  custom_metadata=None
  tree={'layer0': {'bias': ArrayMetadata :  name=layer0.bias,  directory=/tmp/checkpoint-manager-items-1/0/state,  shape=(4,),  sharding=None,  dtype=float64,  storage=StorageMetadata(chunk_shape=(4,), write_shape=None),, 'weight': ArrayMetadata :  name=layer0.weight,  directory=/tmp/checkpoint-manager-items-1/0/state,  shape=(16,),  sharding=SingleDeviceShardingMetadata(device_str=cpu:0),  dtype=int32,  storage=StorageMetadata(chunk_shape=(16,), write_shape=(16,)),}}
  use_zarr3=False
)})
```

--------------------------------

### Validate AsyncOptions and initialize directory

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_manager.html

Ensures that post-finalization callbacks are only used with AsyncCheckpointer and sets up the root directory.

```python
    if (
        self._options.async_options is not None
        and self._options.async_options.post_finalization_callback is not None
        and not isinstance(self._checkpointer, AsyncCheckpointer)
    ):
      raise ValueError(
          'AsyncOptions.post_finalization_callback is only supported with'
          ' AsyncCheckpointer. But final resolved checkpointer is: '
          f' {self._checkpointer}'
      )

    self._directory = epath.Path(directory)
    if self._options.read_only:
      logging.warning('Given directory is read only=%s', self._directory)
    if self._options.create:
      _create_root_directory(
          self._directory,
          self._multiprocessing_options,
          self._options.file_options,
      )
```

--------------------------------

### Checkpoint Composite Output

Source: https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html

Example output structure for a composite checkpoint containing state and metadata.

```text
Composite({'extra_metadata': {'version': 1.0}, 'state': {'a': 0}})
```

--------------------------------

### Define Sharding for SingleReplicaArrayHandler

Source: https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.type_handlers.html

Example of a sharding configuration that must be compatible with the primary replica structure.

```python
shape = (32, 2)
mesh = jax.sharding.Mesh(jax.devices().reshape(shape), ('x', 'y'))
pspec = jax.sharding.PartitionSpec(None, 'y')
sharding=jax.sharding.NamedSharding(mesh, pspec)
```

--------------------------------

### Implement a complete benchmark

Source: https://orbax.readthedocs.io/en/latest/checkpointing_benchmark_guide.html

A full implementation of a benchmark including options definition, decorator binding, and the test function logic.

```python
import dataclasses

import jax
from orbax.checkpoint import v1 as ocp
from orbax.checkpoint._src.testing.benchmarks.core import core as benchmarks_core
from orbax.checkpoint._src.testing.benchmarks.core import metric as metric_lib


# ① Your knobs. A list-typed field is swept (see §8).
@dataclasses.dataclass(frozen=True)
class SaveRestoreOptions(benchmarks_core.BenchmarkOptions):
  use_ocdbt: bool | list[bool] = True

  def is_valid(self) -> bool:        # optional: drop nonsensical combinations
    return True


# ② Bind the options to the generator.
@benchmarks_core.benchmark_options(SaveRestoreOptions)
class SaveRestoreBenchmark(benchmarks_core.BenchmarksGenerator):

  # ③ The measured code. Called once per (option × mesh × checkpoint).
  def test_fn(
      self, context: benchmarks_core.TestContext
  ) -> benchmarks_core.TestResult:
    metrics = metric_lib.Metrics()

    options = context.options                  # a SaveRestoreOptions instance
    pytree = context.pytree                    # generated from checkpoint_config.spec
    abstract = jax.tree.map(ocp.arrays.to_shape_dtype_struct, pytree)
    path = context.path / "ckpt"               # a fresh per-run directory

    with ocp.Context(
        array_options=ocp.options.ArrayOptions(
            saving=ocp.options.ArrayOptions.Saving(use_ocdbt=options.use_ocdbt)
        )
    ):
      with metrics.measure("save"):
        ocp.save(path, pytree)
      with metrics.measure("load"):
        _ = ocp.load(path, abstract_state=abstract)

    return benchmarks_core.TestResult(metrics=metrics)
```

--------------------------------

### Get default preservation policy

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_manager.html

Constructs a default preservation policy based on the provided CheckpointManagerOptions.

```python
def _get_default_preservation_policy(
    options: CheckpointManagerOptions,
) -> preservation_policy_lib.PreservationPolicy:
  """Returns a default preservation policy."""
  preservation_policies = []
  if options.keep_period is not None:
    preservation_policies.append(
        preservation_policy_lib.EveryNSteps(options.keep_period)
    )
  if options.keep_time_interval is not None:
    total_seconds = int(options.keep_time_interval.total_seconds())
    preservation_policies.append(
        preservation_policy_lib.EveryNSeconds(interval_secs=total_seconds)
    )
  if options.best_fn is not None:
    preservation_policies.append(
        preservation_policy_lib.BestN(
            get_metric_fn=options.best_fn,
            reverse=(options.best_mode == 'min'),
            n=options.max_to_keep,
            keep_checkpoints_without_metrics=options.keep_checkpoints_without_metrics,
        )
    )
  else:
    preservation_policies.append(
        preservation_policy_lib.LatestN(n=options.max_to_keep)
    )
  return preservation_policy_lib.AnyPreservationPolicy(preservation_policies)
```

--------------------------------

### Initialize PyTreeCheckpointHandler

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.html

Demonstrates the basic instantiation of a Checkpointer using the PyTreeCheckpointHandler.

```python
ckptr = Checkpointer(PyTreeCheckpointHandler())
```

--------------------------------

### Release Snapshot

Source: https://orbax.readthedocs.io/en/latest/_modules/orbax/checkpoint/checkpoint_utils.html

Releases a checkpoint snapshot by deleting it. Only executed by the primary process in a multi-host setup.

```python
def _release_snapshot(
    checkpoint_dir: epath.Path,
    step: int,
    step_name_format: step_lib.NameFormat[step_lib.Metadata],
    snapshot_dir: Optional[epath.Path] = None,
    *,
    ignore_file_not_found_error: bool | None = None,
):
  """Releases snapshot by deleting the snapshot of the checkpoint."""
  if multihost.process_index() == 0:
    logging.info('Releasing snapshot at step: %d.', step)
    if snapshot_dir is None:
      snapshot_dir = checkpoint_dir / _SNAPSHOTS
    snapshot_path = snapshot_dir / step_name_format.build_name(step)
    snapshot_impl = snapshot_lib.create_instance(checkpoint_dir, snapshot_path)
    with _manage_snapshot_file_not_found(
        ignore_errors=ignore_file_not_found_error,  # pyrefly: ignore[bad-argument-type]
        formatted_message=(
            f'Ignoring error when releasing snapshot for step: {step}'
        ),
    ):
      asyncio_utils.run_sync(snapshot_impl.release_snapshot())
```

--------------------------------

### CheckpointManagerOptions.__init__

Source: https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpoint_manager.html

Initializes the configuration options for a CheckpointManager instance.

```APIDOC
## CheckpointManagerOptions.__init__

### Description
Initializes the configuration options for the CheckpointManager. These options define how checkpoints are saved, which ones are preserved, and how the underlying file system is managed.

### Parameters
- **save_interval_steps** (int) - Optional - Interval at which to save checkpoints.
- **max_to_keep** (int) - Optional - Maximum number of checkpoints to keep.
- **keep_time_interval** (float) - Optional - Time interval for keeping checkpoints.
- **keep_period** (int) - Optional - Period for keeping checkpoints.
- **should_keep_fn** (callable) - Optional - Predicate to determine if a checkpoint should be kept.
- **best_fn** (callable) - Optional - Function to determine the best checkpoint.
- **best_mode** (str) - Optional - Mode for best_fn, default 'max'.
- **keep_checkpoints_without_metrics** (bool) - Optional - Whether to keep checkpoints without metrics, default True.
- **step_prefix** (str) - Optional - Prefix for step names.
- **step_format_fixed_length** (int) - Optional - Fixed length for step formatting.
- **step_name_format** (str) - Optional - Format string for step names.
- **create** (bool) - Optional - Whether to create directories, default True.
- **cleanup_tmp_directories** (bool) - Optional - Whether to clean up temporary directories, default False.
- **save_on_steps** (list) - Optional - Specific steps to save on.
- **single_host_load_and_broadcast** (bool) - Optional - Whether to load and broadcast on a single host, default False.
- **todelete_subdir** (str) - Optional - Subdirectory for deletions.
- **todelete_full_path** (str) - Optional - Full path for deletions.
- **enable_background_delete** (bool) - Optional - Whether to enable background deletion, default False.
- **num_deletion_threads** (int) - Optional - Number of threads for deletion.
- **read_only** (bool) - Optional - Whether to open in read-only mode, default False.
- **enable_async_checkpointing** (bool) - Optional - Whether to enable async checkpointing, default True.
- **async_options** (object) - Optional - Options for async checkpointing.
- **multiprocessing_options** (object) - Optional - Options for multiprocessing.
- **should_save_fn** (callable) - Optional - Predicate to check if a step should be saved.
- **file_options** (FileOptions) - Optional - Configuration for directories and files.
- **save_root_metadata** (bool) - Optional - Whether to save root-level metadata, default True.
- **temporary_path_class** (class) - Optional - Concrete class for temporary paths.
- **save_decision_policy** (object) - Optional - Policy object for save decisions.
- **preservation_policy** (object) - Optional - Policy object for preservation decisions.
- **prevent_write_metrics** (bool) - Optional - Whether to prevent writing metrics, default False.
- **enable_should_save_is_saving_in_progress_check** (bool) - Optional - Whether to check if saving is in progress, default True.
- **enable_per_process_directory_creation** (bool) - Optional - Whether to create directories per process, default False.
- **lightweight_initialize** (bool) - Optional - Whether to perform lightweight initialization, default False.
```
Tokens:
Raw