Checkpoint

substrax.checkpoint is one checkpoint store over Orbax: step-addressed, pickle-free, with retention delegated to orbax.checkpoint.CheckpointManager. A checkpoint is a step holding named items, the things a training loop owns (model, optimizer, rng, data_iterator, extensions), beside one CheckpointMetadata record.

from flax import nnx
from substrax.checkpoint import OrbaxCheckpointStore

with OrbaxCheckpointStore("checkpoints", max_to_keep=5) as store:
    store.save(100, {"model": nnx.state(model), "rng": key}, epoch=2, metrics={"loss": 0.42})
    checkpoint = store.restore(100, templates={"model": nnx.state(model)})
    nnx.update(model, checkpoint.items["model"])
    best = store.best_step("loss")

Every item is written with PyTreeSave, which carries arrays, typed PRNG keys and plain-Python leaves (ints, strings, booleans, lists) alike, and the metadata with JsonSave, so restoring a checkpoint never executes code. On disk each item holds its pytree under a single tree node, so an item that is a bare array (the rng key) is accepted by every supported Orbax, the 0.11.33 floor included. Restoring onto templates places every array on its template leaf's device, whatever topology the checkpoint was written on; an item without a template comes back as stored.

A template never changes a saved array's dtype silently. Orbax casts every array to its template leaf's dtype, so a float64 array restored onto a float32 template, or a float32 model onto a bfloat16 one, would lose precision or change without a word. save writes every array's dtype, by item and leaf path, as a dtypes JSON item beside the metadata record (which best_step reads for every step, so it stays small), and restore compares each template leaf against it before reading any array, raising CheckpointDtypeMismatchError with a DtypeMismatch (item, leaf, saved and restored dtype) for every differing array. An item restored without a template is compared after it is read: while jax's x64 mode is off, jax creates a 64-bit array at 32 bits, and the error says how to enable x64. restore(..., cast_dtypes=True) accepts the cast in either case. Template leaves may be arrays or jax.ShapeDtypeStruct; a typed PRNG key compares by its uint32 data, as Orbax stores it, and strings and Python numbers are not compared. A checkpoint written without the dtypes item is compared against Orbax's per-array metadata instead, which opens every array's store on each restore.

Writes are strict: an existing step is refused unless overwrite=True, a step below the latest is refused, and a step Orbax declines raises. Every refusal is a CheckpointNotWrittenError naming the step, the latest step and the reason. Nothing is created before the first save, and resolve_checkpoint_dir picks a run's directory without creating it.

A checkpoint whose metadata names another format, another producer or another version number is refused with UnsupportedCheckpointError naming what it holds, rather than read as this one.

substrax.checkpoint

Checkpoints: named items beside one metadata record.

A checkpoint is a step holding the items a training loop owns (model, optimizer, rng, data_iterator, extensions), each an Orbax pytree item, and one :class:CheckpointMetadata record. :class:OrbaxCheckpointStore is the one implementation of the :class:CheckpointStore protocol; it refuses to overwrite or reorder steps unless told to, restores onto templates so a checkpoint written on one device topology lands on whatever the restoring process has, and records each array's dtype at save so a restore that would change one is refused rather than silent.

CheckpointDtypeMismatchError(*, step, mismatches, x64_disabled=False)

Bases: ValueError

A restore would give saved arrays another dtype, so it is refused.

Orbax casts every array to its template leaf's dtype, and without a template jax creates a 64-bit array at 32 bits while its x64 mode is off. step is the checkpoint's step and mismatches every array concerned; restore(..., cast_dtypes=True) accepts the cast.

Record the step and every array concerned.

Parameters:

Name	Type	Description	Default
step	int	The checkpoint's step.	required
mismatches	tuple[DtypeMismatch, ...]	Every array whose restored dtype differs from its saved one.	required
x64_disabled	bool	Whether jax's x64 mode was off for a restore without templates, which adds how to keep the saved precision to the message.	False

CheckpointNotFoundError(step, directory)

Bases: FileNotFoundError

The store holds no checkpoint at the requested step.

Record the missing step and the store's directory.

CheckpointNotWrittenError(*, step, latest_step, reason)

Bases: RuntimeError

A save was refused, so the step on disk is whatever it was before the call.

The error carries step, the step the caller asked to write; latest_step, the newest step the store held at the time (None for an empty store); and reason: "exists" when the step is already written and overwrite was not given, "below_latest" when the step is older than the newest one (a store's steps only grow), "rejected" when Orbax declined the write.

Record the refused step, the store's latest step and the reason.

DtypeMismatch(*, item, leaf, saved, restored) dataclass

One array a restore would give another dtype than the one it was saved with.

item is the Orbax item on disk and leaf the array's path inside it joined with /. saved is the dtype written; restored the dtype the restore gives it: the template leaf's, or, without a template, the one this process creates for the saved array.

UnsupportedCheckpointError

Bases: ValueError

The checkpoint's metadata describes a format this substrax cannot read.

CheckpointMetadata(*, format=FORMAT_NAME, format_version=CURRENT_FORMAT_VERSION, step, epoch=None, items, libraries=dict(), producer=None, metrics=dict(), extra=dict(), created_at) dataclass

What a checkpoint says about itself.

Attributes:

Name	Type	Description
format	str	"substrax-checkpoint"; another name is another producer's file.
format_version	int	The on-disk format; a reader refuses a version newer than its own.
step	int	The step the checkpoint is addressed by.
epoch	int | None	The epoch the step falls in, when the producer counts epochs.
items	tuple[str, ...]	The item names written, in the order they were given.
libraries	Mapping[str, str]	Distribution versions at write time, informational only.
producer	Producer | None	The library that wrote the checkpoint, when it says.
metrics	Mapping[str, float]	Scalar metrics, the values best_step compares.
extra	Mapping[str, JsonValue]	The caller's own JSON values; a key here never shadows a field of the record.
created_at	str	The write time, ISO 8601 in UTC.

to_dict()

The JSON form Orbax writes as the metadata item.

from_dict(payload) classmethod

Read the JSON form, refusing a record this substrax cannot interpret.

Parameters:

Name	Type	Description	Default
payload	Mapping[str, JsonValue]	The record as Orbax's JsonRestore returns it.	required

Returns:

Type	Description
CheckpointMetadata	The metadata record.

Raises:

Type	Description
UnsupportedCheckpointError	If the record names another format or another version of it.
ValidationError	If a field is missing or holds a value its annotation does not admit.

Producer(*, name, version) dataclass

The library that wrote a checkpoint, by distribution name and version.

to_dict()

The JSON form.

from_dict(payload) classmethod

Read the JSON form.

Parameters:

Name	Type	Description	Default
payload	Mapping[str, JsonValue]	The producer's JSON object.	required

Returns:

Type	Description
Producer	The producer.

Raises:

Type	Description
ValidationError	If name or version is missing or not a string.

Checkpoint(*, step, items, metadata) dataclass

A restored checkpoint: its step, its items by name and its metadata.

CheckpointStore

Bases: Protocol

Step-addressed checkpoint store contract.

Implementations persist named items plus a metadata record under an integer step and restore them back. Training code depends on this protocol; the Orbax class is the one implementation.

save(step, items, *, epoch=None, metrics=None, producer=None, extra=None, overwrite=False)

Write items at step and return the checkpoint's directory.

restore(step, *, templates=None, cast_dtypes=False)

Read the checkpoint at step, onto templates where given.

read_metadata(step)

Read only the metadata record of step.

list_steps()

Every step held, ascending.

latest_step()

The newest step, or None for an empty store.

best_step(metric, *, mode='min')

The step whose metric is the lowest or highest, or None when none has it.

delete(step)

Remove the checkpoint at step.

close()

Release backend resources.

OrbaxCheckpointStore(directory, *, max_to_keep=5)

Orbax-backed :class:CheckpointStore.

Wraps :class:orbax.checkpoint.CheckpointManager with step-int addressing and retention by max_to_keep. The manager is opened on first use, so constructing a store writes nothing; the directory appears with the first save. The store doubles as a context manager so backend resources are released deterministically::

with OrbaxCheckpointStore(path, max_to_keep=5) as store:
    store.save(step, {"model": nnx.state(model)}, metrics={"loss": loss})

Remember where the checkpoints live and how many to keep.

Parameters:

Name	Type	Description	Default
directory	str | Path	Directory the checkpoints are stored under; created on the first save.	required
max_to_keep	int | None	How many of the newest steps Orbax retains; None keeps every one.	5

Raises:

Type	Description
ValueError	If directory is empty or whitespace-only.

__enter__()

Enter the context manager.

__exit__(*_exc_info)

Close the backend on context exit.

save(step, items, *, epoch=None, metrics=None, producer=None, extra=None, overwrite=False)

Write items at step with a metadata record and every array's dtype.

Parameters:

Name	Type	Description	Default
step	int	Non-negative step the checkpoint is addressed by.	required
items	Mapping[str, PyTree]	The pytrees to write, by item name (one of :data:~substrax.checkpoint.metadata.ITEM_NAMES).	required
epoch	int | None	The epoch the step falls in, when the caller counts epochs.	None
metrics	Mapping[str, float] | None	Scalar metrics recorded for best_step.	None
producer	Producer | None	The library writing the checkpoint.	None
extra	Mapping[str, JsonValue] | None	The caller's own JSON values, under keys that are not fields of the record.	None
overwrite	bool	Replace an existing checkpoint at step instead of refusing.	False

Returns:

Type	Description
Path	The checkpoint's directory, directory / str(step).

Raises:

Type	Description
ValueError	If step is negative, items is empty or names an unknown item, or extra shadows a reserved key.
CheckpointNotWrittenError	If step is already written and overwrite is false, if it is below the latest step, or if Orbax declines the write.

restore(step, *, templates=None, cast_dtypes=False)

Read the checkpoint at step.

Parameters:

Name	Type	Description	Default
step	int	The step to read.	required
templates	Mapping[str, PyTree] | None	Pytrees by item name; each named item is restored onto its template's leaves (device placement included), and an item without a template comes back as stored. Leaves may be arrays or jax.ShapeDtypeStruct.	None
cast_dtypes	bool	Accept an array coming back with another dtype than it was saved with: cast to its template leaf's, or, without a template, a 64-bit array at 32 bits while jax's x64 mode is off. Without it, either is refused.	False

Returns:

Type	Description
Checkpoint	The restored items and the checkpoint's metadata.

Raises:

Type	Description
CheckpointNotFoundError	If the store holds no checkpoint at step.
UnsupportedCheckpointError	If the checkpoint is in another format than the one this substrax reads, or is not a substrax checkpoint.
CheckpointDtypeMismatchError	If an array would come back with another dtype than it was saved with and cast_dtypes is false.
ValueError	If templates names an item the checkpoint lacks, or a template's tree does not match the checkpoint's.

read_metadata(step)

Read the metadata record of step without its items.

Parameters:

Name	Type	Description	Default
step	int	The step to read.	required

Returns:

Type	Description
CheckpointMetadata	The metadata record.

Raises:

Type	Description
CheckpointNotFoundError	If the store holds no checkpoint at step.
UnsupportedCheckpointError	If the checkpoint is in another format than the one this substrax reads, or is not a substrax checkpoint.

list_steps()

Every step held, ascending.

latest_step()

The newest step, or None for an empty store.

best_step(metric, *, mode='min')

The step whose metric is the lowest ("min") or highest ("max").

Steps whose metadata lacks the metric are skipped.

Parameters:

Name	Type	Description	Default
metric	str	The metric name, as recorded in each step's metrics.	required
mode	BestMode	"min" for the lowest value, "max" for the highest.	'min'

Returns:

Type	Description
int | None	The best step, or None when no step records the metric.

Raises:

Type	Description
ValueError	If mode is neither "min" nor "max".

delete(step)

Remove the checkpoint at step.

Parameters:

Name	Type	Description	Default
step	int	The step to remove.	required

Raises:

Type	Description
CheckpointNotFoundError	If the store holds no checkpoint at step.

close()

Close the Orbax manager, if one was opened.

resolve_checkpoint_dir(checkpoint_dir, run_dir)

Return the directory a run's checkpoints go to, without creating it.

An explicit checkpoint_dir wins; otherwise the run's directory holds a checkpoints subdirectory. A relative path resolves against the working directory. The store creates the directory on its first save, so resolving one writes nothing.

Parameters:

Name	Type	Description	Default
checkpoint_dir	Path | None	The directory the caller chose, or None.	required
run_dir	Path | None	The run's directory, whose checkpoints subdirectory is the default.	required

Returns:

Type	Description
Path	The absolute directory the checkpoints go to.

Raises:

Type	Description
ValueError	If neither directory is given.