# {py:mod}`evorl.utils.orbax_utils`

```{py:module} evorl.utils.orbax_utils
```

```{autodoc2-docstring} evorl.utils.orbax_utils
:parser: autodoc2_docstrings_parser
:allowtitles:
```

## Module Contents

### Classes

````{list-table}
:class: autosummary longtable
:align: left

* - {py:obj}`CheckpointManager <evorl.utils.orbax_utils.CheckpointManager>`
  -
* - {py:obj}`DummyCheckpointManager <evorl.utils.orbax_utils.DummyCheckpointManager>`
  - ```{autodoc2-docstring} evorl.utils.orbax_utils.DummyCheckpointManager
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### Functions

````{list-table}
:class: autosummary longtable
:align: left

* - {py:obj}`filter_zero_size_arrays_on_restore <evorl.utils.orbax_utils.filter_zero_size_arrays_on_restore>`
  - ```{autodoc2-docstring} evorl.utils.orbax_utils.filter_zero_size_arrays_on_restore
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`filter_zero_size_arrays_on_save <evorl.utils.orbax_utils.filter_zero_size_arrays_on_save>`
  - ```{autodoc2-docstring} evorl.utils.orbax_utils.filter_zero_size_arrays_on_save
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`load <evorl.utils.orbax_utils.load>`
  - ```{autodoc2-docstring} evorl.utils.orbax_utils.load
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`save <evorl.utils.orbax_utils.save>`
  - ```{autodoc2-docstring} evorl.utils.orbax_utils.save
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`setup_checkpoint_manager <evorl.utils.orbax_utils.setup_checkpoint_manager>`
  - ```{autodoc2-docstring} evorl.utils.orbax_utils.setup_checkpoint_manager
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

`````{py:class} CheckpointManager(directory: etils.epath.PathLike, checkpointers: typing.Optional[typing.Union[orbax.checkpoint.checkpoint_manager.AbstractCheckpointer, orbax.checkpoint.checkpoint_manager.CheckpointersDict]] = None, options: typing.Optional[orbax.checkpoint.checkpoint_manager.CheckpointManagerOptions] = None, metadata: typing.Optional[typing.Mapping[str, typing.Any]] = None, item_names: typing.Optional[typing.Sequence[str]] = None, item_handlers: typing.Optional[typing.Union[orbax.checkpoint.checkpoint_manager.CheckpointHandler, orbax.checkpoint.checkpoint_manager.CheckpointHandlersDict]] = None, logger: typing.Optional[orbax.checkpoint._src.logging.abstract_logger.AbstractLogger] = None, handler_registry: typing.Optional[orbax.checkpoint.checkpoint_manager.CheckpointHandlerRegistry] = None)
:canonical: evorl.utils.orbax_utils.CheckpointManager

Bases: {py:obj}`orbax.checkpoint.CheckpointManager`

````{py:method} restore(step, items, **kwargs)
:canonical: evorl.utils.orbax_utils.CheckpointManager.restore

````

````{py:method} save(step, items, **kwargs)
:canonical: evorl.utils.orbax_utils.CheckpointManager.save

````

`````

`````{py:class} DummyCheckpointManager
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager

Bases: {py:obj}`orbax.checkpoint.AbstractCheckpointManager`

```{autodoc2-docstring} evorl.utils.orbax_utils.DummyCheckpointManager
:parser: autodoc2_docstrings_parser
```

````{py:method} all_steps(read: bool = False) -> collections.abc.Sequence[int]
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.all_steps

````

````{py:method} best_step() -> int | None
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.best_step

````

````{py:method} check_for_errors()
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.check_for_errors

````

````{py:method} close()
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.close

````

````{py:method} delete(step: int)
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.delete

````

````{py:method} directory()
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.directory

````

````{py:method} item_metadata(step: int)
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.item_metadata

````

````{py:method} latest_step() -> int | None
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.latest_step

````

````{py:method} metadata() -> collections.abc.Mapping[str, typing.Any]
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.metadata

````

````{py:method} metrics(step: int) -> typing.Any | None
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.metrics

````

````{py:method} reached_preemption(step: int) -> bool
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.reached_preemption

````

````{py:method} reload()
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.reload

````

````{py:method} restore(step, items, **kwargs)
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.restore
:abstractmethod:

````

````{py:method} save(step, items, **kwargs)
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.save

````

````{py:method} should_save(step: int) -> bool
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.should_save

````

````{py:method} wait_until_finished()
:canonical: evorl.utils.orbax_utils.DummyCheckpointManager.wait_until_finished

````

`````

````{py:function} filter_zero_size_arrays_on_restore(src_tree: chex.ArrayTree, dst_tree: chex.ArrayTree) -> chex.ArrayTree
:canonical: evorl.utils.orbax_utils.filter_zero_size_arrays_on_restore

```{autodoc2-docstring} evorl.utils.orbax_utils.filter_zero_size_arrays_on_restore
:parser: autodoc2_docstrings_parser
```
````

````{py:function} filter_zero_size_arrays_on_save(tree: chex.ArrayTree) -> chex.ArrayTree
:canonical: evorl.utils.orbax_utils.filter_zero_size_arrays_on_save

```{autodoc2-docstring} evorl.utils.orbax_utils.filter_zero_size_arrays_on_save
:parser: autodoc2_docstrings_parser
```
````

````{py:function} load(path, state: chex.ArrayTree) -> chex.ArrayTree
:canonical: evorl.utils.orbax_utils.load

```{autodoc2-docstring} evorl.utils.orbax_utils.load
:parser: autodoc2_docstrings_parser
```
````

````{py:function} save(path, state: chex.ArrayTree)
:canonical: evorl.utils.orbax_utils.save

```{autodoc2-docstring} evorl.utils.orbax_utils.save
:parser: autodoc2_docstrings_parser
```
````

````{py:function} setup_checkpoint_manager(config: omegaconf.DictConfig) -> orbax.checkpoint.CheckpointManager
:canonical: evorl.utils.orbax_utils.setup_checkpoint_manager

```{autodoc2-docstring} evorl.utils.orbax_utils.setup_checkpoint_manager
:parser: autodoc2_docstrings_parser
```
````