# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0

"""Central components used by other modules.

Central parts of composer such as :class:`~.engine.Engine`, base class for critical components such as
:class:`~.algorithm.Algorithm` and :class:`~.callback.Callback` and other useful functionality such as
:class:`~.time.Timestamp` are implemented under core.
"""

from composer.core.algorithm import Algorithm
from composer.core.callback import Callback
from composer.core.data_spec import DataSpec, ensure_data_spec
from composer.core.engine import Engine, Trace
from composer.core.evaluator import Evaluator, ensure_evaluator
from composer.core.event import Event
from composer.core.passes import AlgorithmPass
from composer.core.precision import Precision, get_precision_context
from composer.core.serializable import Serializable
from composer.core.state import State
from composer.core.time import Time, Timestamp, TimeUnit, ensure_time
from composer.core.types import JSON, Batch, BreakEpochException, Dataset, MemoryFormat, PyTorchScheduler, TrainerMode

__all__ = [
    'Algorithm',
    'AlgorithmPass',
    'Callback',
    'DataSpec',
    'Dataset',
    'ensure_data_spec',
    'Engine',
    'Trace',
    'Evaluator',
    'Event',
    'Precision',
    'get_precision_context',
    'Serializable',
    'State',
    'Time',
    'Timestamp',
    'TimeUnit',
    'ensure_time',
    'ensure_evaluator',
    'Batch',
    'PyTorchScheduler',
    'JSON',
    'MemoryFormat',
    'TrainerMode',
    'BreakEpochException',
]
