kd.ckpts.items.CheckpointItem

class kauldron.checkpoints.checkpoint_items.CheckpointItem

Bases: abc.ABC

Interface for a checkpoint item.

This interface defines how instances should be saved by orbax.

Protocol is equivalent to:

item = MyCheckpointItem(values=...)  # Subclass of `CheckpointItem`

ckpt = ocp.CheckpointManager(item_handlers=item.__kd_ocp_handlers__())
ckpt.save(step, args=item.__kd_ocp_save_args__())

out = ckpt.restore(step, args=item.__kd_ocp_restore_args__())
restored_item = item.__kd_ocp_restore_post__(out)