src/checkpointing/__init__.py

"""
Pico Checkpointing Package

We subdivide the checkpointing into training, evaluation, and learning_dynamics. Training
checkpoints store the model, optimizer, and learning rate scheduler. Evaluation checkpoints store
the evaluation results on the defined metrics. Learning dynamics checkpoints store activations and gradients used for
learning dynamics analysis.
"""

from .evaluation import save_evaluation_results
from .learning_dynamics import (
    compute_learning_dynamics_states,
    save_learning_dynamics_states,
)
from .training import load_checkpoint, save_checkpoint

__all__ = [
    "compute_learning_dynamics_states",
    "load_checkpoint",
    "save_checkpoint",
    "save_evaluation_results",
    "save_learning_dynamics_states",
]