"""Module implementing the Feed-Forward Neural Network (FFN) architectures."""

from llmfoundry.layers_registry import (
    module_init_fns,
)
from llmfoundry.registry import ffns

from mt_dao.models.ffns.builder import build_sigma_moe, sigma_moe_init_sp
from mt_dao.models.ffns.sigma_moe import (
    SigmaMoE,
    SigmaMoEConfig,
)
from mt_dao.models.ffns.utils import entropy_l, entropy_reg, log_mean

ffns.register("sigma_moe", func=build_sigma_moe)

module_init_fns.register(
    "sigma_moe",
    func=sigma_moe_init_sp,
)

__all__ = [
    "SigmaMoE",
    "SigmaMoEConfig",
    "entropy_l",
    "entropy_reg",
    "log_mean",
]
