Checkpointing¶

tinax.checkpointing owns Orbax V1 checkpointables and explicit abstract restoration targets. It atomically saves related resume-critical values together and returns asynchronous responses for callers to complete.

See the Training Checkpoints guide for task-oriented usage.

save_checkpointables ¶

save_checkpointables(
    path: str | PathLike[str],
    checkpointables: Mapping[str, Any],
) -> AsyncResponse[None]

Start one immutable atomic V1 save on every process's main thread and return its completion response.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination checkpoint directory. Must not use Orbax's reserved temporary suffix.	required
`checkpointables`	`Mapping[str, Any]`	Non-empty mapping of checkpointable name to value. Names must be portable and distinct.	required

Returns:

Type	Description
`AsyncResponse[None]`	An Orbax `AsyncResponse`; call `.result()` to await completion before
`AsyncResponse[None]`	loading from or deleting the destination.

Raises:

Type	Description
`TypeError`	If `checkpointables` is not a mapping or a name is not a string.
`ValueError`	If `checkpointables` is empty, a name is invalid or not distinct, or `path` uses the reserved temporary suffix.

load_checkpointables ¶

load_checkpointables(
    path: str | PathLike[str], targets: Mapping[str, Any]
) -> dict[str, Any]

Load all named V1 targets together on every process's main thread after save completion.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Source checkpoint directory.	required
`targets`	`Mapping[str, Any]`	Non-empty mapping of checkpointable name to an explicit non-`None` restore target. Names must be portable and distinct.	required

Returns:

Type	Description
`dict[str, Any]`	A dict mapping each requested name to its restored value.

Raises:

Type	Description
`TypeError`	If `targets` is not a mapping or a name is not a string.
`ValueError`	If `targets` is empty, a name is invalid or not distinct, a target is `None`, or `path` uses the reserved temporary suffix.

abstract_restore_target ¶

abstract_restore_target(state: Any) -> Any

Replace JAX and NumPy arrays with shape/dtype targets while preserving JAX sharding and static leaves.

Parameters:

Name	Type	Description	Default
`state`	`Any`	Arbitrary pytree of restore targets. JAX arrays become `jax.ShapeDtypeStruct` (retaining sharding); NumPy arrays become an internal shape/dtype target; other leaves are returned unchanged.	required

Returns:

Type	Description
`Any`	A pytree matching `state` with concrete arrays replaced by abstract
`Any`	shape/dtype restore targets.

TrainingCheckpoint `dataclass` ¶

TrainingCheckpoint(
    step: int,
    model: Any,
    optimizer: Any,
    rng: Any,
    auxiliary: Any,
    iterator: DatasetIterator[Any],
)

A complete restored training state from one atomic checkpoint.

Attributes:

Name	Type	Description
`step`	`int`	Non-negative training step the checkpoint was written at.
`model`	`Any`	Restored model state.
`optimizer`	`Any`	Restored optimizer state.
`rng`	`Any`	Restored RNG state.
`auxiliary`	`Any`	Restored user auxiliary state.
`iterator`	`DatasetIterator[Any]`	Restored Grain input iterator.

Raises:

Type	Description
`TypeError`	If `step` is not an integer (booleans rejected) or `iterator` is not a `grain.DatasetIterator`.
`ValueError`	If `step` is negative.

TrainingCheckpointNames `dataclass` ¶

TrainingCheckpointNames(
    model: str = "model",
    optimizer: str = "optimizer",
    rng: str = "rng",
    auxiliary: str = "auxiliary",
    iterator: str = "iterator",
)

Names for the five independently handled parts of a training checkpoint.

Attributes:

Name	Type	Description
`model`	`str`	Checkpointable name for model state.
`optimizer`	`str`	Checkpointable name for optimizer state.
`rng`	`str`	Checkpointable name for RNG state.
`auxiliary`	`str`	Checkpointable name for auxiliary state and step metadata.
`iterator`	`str`	Checkpointable name for the input iterator state.

Raises:

Type	Description
`TypeError`	If any name is not a string.
`ValueError`	If any name is invalid (see `validate_checkpointable_name`) or the names are not portable and distinct.

save_training_checkpoint ¶

save_training_checkpoint(
    path: str | PathLike[str],
    step: int,
    model: Any,
    optimizer: Any,
    rng: Any,
    auxiliary: Any,
    iterator: DatasetIterator[Any],
    *,
    names: TrainingCheckpointNames | None = None,
) -> AsyncResponse[None]

Start one atomic training save on every process's main thread and return its completion response.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination checkpoint directory.	required
`step`	`int`	Non-negative training step.	required
`model`	`Any`	Model state to save.	required
`optimizer`	`Any`	Optimizer state to save.	required
`rng`	`Any`	RNG state to save.	required
`auxiliary`	`Any`	User auxiliary state saved alongside `step`.	required
`iterator`	`DatasetIterator[Any]`	Grain input iterator whose state is saved.	required
`names`	`TrainingCheckpointNames \| None`	Optional override for the five checkpointable names. Defaults to `TrainingCheckpointNames()`.	`None`

Returns:

Type	Description
`AsyncResponse[None]`	An Orbax `AsyncResponse`; call `.result()` to await completion.

Raises:

Type	Description
`TypeError`	If `step` is not an integer (booleans rejected), `iterator` is not a `grain.DatasetIterator`, or `names` is not a `TrainingCheckpointNames`.
`ValueError`	If `step` is negative.

load_training_checkpoint ¶

load_training_checkpoint(
    path: str | PathLike[str],
    model_target: Any,
    optimizer_target: Any,
    rng_target: Any,
    auxiliary_target: Any,
    iterator: DatasetIterator[Any],
    *,
    names: TrainingCheckpointNames | None = None,
) -> TrainingCheckpoint

Preflight value targets before restoring the stateful iterator from the same atomic checkpoint.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Source checkpoint directory.	required
`model_target`	`Any`	Restore target for model state.	required
`optimizer_target`	`Any`	Restore target for optimizer state.	required
`rng_target`	`Any`	Restore target for RNG state.	required
`auxiliary_target`	`Any`	Restore target for user auxiliary state.	required
`iterator`	`DatasetIterator[Any]`	Grain iterator to restore in place.	required
`names`	`TrainingCheckpointNames \| None`	Optional override for the five checkpointable names. Defaults to `TrainingCheckpointNames()`.	`None`

Returns:

Type	Description
`TrainingCheckpoint`	A `TrainingCheckpoint` with the restored step, model, optimizer, rng,
`TrainingCheckpoint`	auxiliary state, and iterator.

Raises:

Type	Description
`TypeError`	If `iterator` is not a `grain.DatasetIterator`, `names` is not a `TrainingCheckpointNames`, or the auxiliary checkpointable does not restore as a mapping.
`ValueError`	If the auxiliary checkpointable is missing `step` or `state`, or the restored step is negative.

validate_checkpointable_name ¶

validate_checkpointable_name(name: object) -> str

Validate one checkpointable name as a safe portable path component.

Parameters:

Name	Type	Description	Default
`name`	`object`	Candidate checkpointable name to validate.	required

Returns:

Type	Description
`str`	The validated `name`, unchanged.

Raises:

Type	Description
`TypeError`	If `name` is not a string.
`ValueError`	If `name` is empty, a dot component, contains reserved or control characters, ends with a space or dot, exceeds 255 UTF-8 bytes, or matches a reserved name.

Legacy (Orbax V0)¶

Warning

The tinax.checkpointing.legacy.v0 namespace is retained only for explicit Orbax V0 compatibility work. Do not use it for new checkpoints.

save_legacy_v0_pytree ¶

save_legacy_v0_pytree(
    path: str | PathLike[str], state: Any
) -> None

Write one immutable legacy V0 PyTree checkpoint synchronously on every process's main thread.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination checkpoint directory. Must not use Orbax's reserved temporary suffix.	required
`state`	`Any`	PyTree of values to checkpoint.	required

Raises:

Type	Description
`ValueError`	If `path` uses Orbax's reserved temporary suffix.

load_legacy_v0_pytree ¶

load_legacy_v0_pytree(
    path: str | PathLike[str], target: Any | None = None
) -> Any

Read one legacy V0 PyTree checkpoint synchronously on every process's main thread.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Source checkpoint directory.	required
`target`	`Any \| None`	Optional PyTree restore target guiding structure and sharding. `None` restores the stored structure.	`None`

Returns:

Type	Description
`Any`	The restored PyTree.

Raises:

Type	Description
`ValueError`	If `path` uses Orbax's reserved temporary suffix.

save_legacy_v0_grain_iterator ¶

save_legacy_v0_grain_iterator(
    path: str | PathLike[str],
    iterator: DatasetIterator[Any],
) -> None

Write one immutable Grain iterator with its legacy V0 handler on every process's main thread.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination checkpoint directory.	required
`iterator`	`DatasetIterator[Any]`	Grain `DatasetIterator` whose state is saved.	required

Raises:

Type	Description
`TypeError`	If `iterator` is not a `grain.DatasetIterator`.
`ValueError`	If `path` uses Orbax's reserved temporary suffix.

load_legacy_v0_grain_iterator ¶

load_legacy_v0_grain_iterator(
    path: str | PathLike[str],
    iterator: DatasetIterator[Any],
) -> DatasetIterator[Any]

Read one Grain iterator with its legacy V0 handler on every process's main thread.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Source checkpoint directory.	required
`iterator`	`DatasetIterator[Any]`	Grain `DatasetIterator` to restore in place.	required

Returns:

Type	Description
`DatasetIterator[Any]`	The restored `grain.DatasetIterator`.

Raises:

Type	Description
`TypeError`	If `iterator` is not a `grain.DatasetIterator`, or the checkpoint did not restore a `DatasetIterator`.
`ValueError`	If `path` uses Orbax's reserved temporary suffix.

Checkpointing¶

save_checkpointables ¶

load_checkpointables ¶

abstract_restore_target ¶

TrainingCheckpoint dataclass ¶

TrainingCheckpointNames dataclass ¶

save_training_checkpoint ¶

load_training_checkpoint ¶

validate_checkpointable_name ¶

Legacy (Orbax V0)¶

save_legacy_v0_pytree ¶

load_legacy_v0_pytree ¶

save_legacy_v0_grain_iterator ¶

load_legacy_v0_grain_iterator ¶

TrainingCheckpoint `dataclass` ¶

TrainingCheckpointNames `dataclass` ¶