kd.ckpts.items.TopLevelCheckpointItem

class kauldron.checkpoints.checkpoint_items.TopLevelCheckpointItem

Bases: kauldron.checkpoints.checkpoint_items.CheckpointItem

Checkpoint item that contains other sub-checkpoint items.

Usage:

class CheckpointState(typing.NamedTuple, TopLevelCheckpointItem):
  train_state: CheckpointItem
  metadata: CheckpointItem
  ds: CheckpointItem

Usually it would be good practice to make this a dataclass(kw_only=True). However in this case, it’s convenient to be able to use it like a NamedTuple, like:

state, timer = ckptr.restore(CheckpointState(state, timer))

DEFAULT_ITEM

Default orbax item (restored if nothing is specified).

Type:

ClassVar[str | None]

DEFAULT_ITEM: ClassVar[str | None] = None