src/manifold/training/__init__.py

from manifold.training.curriculum import (
    CurriculumScheduler,
    CurriculumStage,
    CURRICULUM_STAGES,
    get_cheater_types_for_difficulty,
)
from manifold.training.losses import (
    kl_dirichlet,
    evidential_loss,
    reconstruction_loss,
    compute_total_loss,
    LossConfig,
)
from manifold.training.callbacks import (
    Callback,
    CheckpointCallback,
    EarlyStoppingCallback,
    WandBCallback,
    ProgressCallback,
    CallbackManager,
)
from manifold.training.trainer import (
    train_epoch,
    validate,
    MANIFOLDTrainer,
)

__all__ = [
    "CurriculumScheduler",
    "CurriculumStage",
    "CURRICULUM_STAGES",
    "get_cheater_types_for_difficulty",
    "kl_dirichlet",
    "evidential_loss",
    "reconstruction_loss",
    "compute_total_loss",
    "LossConfig",
    "Callback",
    "CheckpointCallback",
    "EarlyStoppingCallback",
    "WandBCallback",
    "ProgressCallback",
    "CallbackManager",
    "train_epoch",
    "validate",
    "MANIFOLDTrainer",
]