"""Structured configuration for the initialization of models."""

from dataclasses import dataclass
from typing import Any


@dataclass
class ModelInitConfig:
    """Configuration for model initialization.

    Attributes
    ----------
    name: str
        The parameter initialization scheme to use.
        Options: 'default_','baseline_',
                    'kaiming_uniform_', 'kaiming_normal_',
                    'neox_init_', 'small_init_',
                    'xavier_uniform_', or
                    'xavier_normal_'.
                    These mimic the parameter initialization methods in PyTorch.
    init_div_is_residual: int|float|str|bool
        Value to divide initial weights by if ``module._is_residual`` is True.
    emb_init_std: float | None
        The standard deviation of the normal distribution
        used to initialize the embedding layer.
    emb_init_uniform_lim: tuple[float, float] | float | None
        The lower and upper limits of the uniform distribution
        used to initialize the embedding layer.
        Mutually exclusive with ``emb_init_std``.
    init_std: float
        The standard deviation of the normal distribution
        used to initialize the model,
        if using the baseline_ parameter initialization scheme.
    init_gain: float
        The gain to use for parameter initialization with
        kaiming or xavier initialization schemes.
    fan_mode: str
        The fan mode to use for parameter initialization with
        kaiming initialization schemes.
    init_nonlinearity: str
        The nonlinearity to use for parameter initialization
        with kaiming initialization schemes.
    See llmfoundry.models.utils.param_init_fns.py
    for info on other param init config options.

    """

    def __init__(  # noqa: PLR0913
        self,
        *_args: Any,  # noqa: ANN401
        name: str,
        fan_mode: str,
        init_nonlinearity: str,
        emb_init_std: float | None,
        emb_init_uniform_lim: float | None,
        init_std: float | None,
        init_gain: float,
        init_div_is_residual: bool,
        **_kwargs: Any,  # noqa: ANN401
    ) -> None:
        """Initialize the ModelInitConfig with the given parameters.

        # NOTE: The defaults will come from the MosaicML library, which provides
        # standard initialization values.

        Parameters
        ----------
        _args: Any
            Positional arguments for future extensibility.
        name: str
            The parameter initialization scheme to use.
            Options: 'default_','baseline_',
                        'kaiming_uniform_', 'kaiming_normal_',
                        'neox_init_', 'small_init_',
                        'xavier_uniform_', or
                        'xavier_normal_'.
                        These mimic the parameter initialization methods in PyTorch.
        init_div_is_residual: int|float|str|bool
            Value to divide initial weights by if ``module._is_residual`` is True.
        emb_init_std: float | None
            The standard deviation of the normal distribution
            used to initialize the embedding layer.
        emb_init_uniform_lim: tuple[float, float] | float | None
            The lower and upper limits of the uniform distribution
            used to initialize the embedding layer.
            Mutually exclusive with ``emb_init_std``.
        init_std: float | None
            The standard deviation of the normal distribution
            used to initialize the model,
            if using the baseline_ parameter initialization scheme.
        init_gain: float
            The gain to use for parameter initialization with
            kaiming or xavier initialization schemes.
        fan_mode: str
            The fan mode to use for parameter initialization with
            kaiming initialization schemes.
        init_nonlinearity: str
            The nonlinearity to use for parameter initialization
            with kaiming initialization schemes.
        **_kwargs: Any
            Additional keyword arguments for future extensibility.
        See llmfoundry.models.utils.param_init_fns.py
        for info on other param init config options.

        """
        self.name = name
        self.fan_mode = fan_mode
        self.init_nonlinearity = init_nonlinearity
        self.init_div_is_residual = init_div_is_residual
        self.emb_init_std = emb_init_std
        self.emb_init_uniform_lim = emb_init_uniform_lim
        self.init_std = init_std
        self.init_gain = init_gain
