def __init__(
        self,
        input_dim: int,
        output_dim: int,
        past_cov_dim: int,
        future_cov_dim: int,
        static_cov_dim: int,
        nr_params: int,
        hidden_size: int,
        ff_size: int,
        num_blocks: int,
        activation: str,
        dropout: float,
        norm_type: Union[str, nn.Module],
        normalize_before: bool,
        **kwargs,
    ) -> None:
        """
        Initializes the TSMixer module for use within a Darts forecasting model.

        Parameters
        ----------
        input_dim
            Number of input target features.
        output_dim
            Number of output target features.
        past_cov_dim
            Number of past covariate features.
        future_cov_dim
            Number of future covariate features.
        static_cov_dim
            Number of static covariate features (number of target features
            (or 1 if global static covariates) * number of static covariate features).
        nr_params
            The number of parameters of the likelihood (or 1 if no likelihood is used).
        hidden_size
           Hidden state size of the TSMixer.
        ff_size
            Dimension of the feedforward network internal to the module.
        num_blocks
            Number of mixer blocks.
        activation
            Activation function to use.
        dropout
            Dropout rate for regularization.
        norm_type
            Type of normalization to use.
        normalize_before
            Whether to apply normalization before or after mixing.
        """
        super().__init__(**kwargs)
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.future_cov_dim = future_cov_dim
        self.static_cov_dim = static_cov_dim
        self.nr_params = nr_params

        if activation not in ACTIVATIONS:
            raise_log(
                ValueError(
                    f"Invalid `activation={activation}`. Must be on of {ACTIVATIONS}."
                ),
                logger=logger,
            )
        activation = getattr(nn, activation)()

        if isinstance(norm_type, str):
            if norm_type not in NORMS:
                raise_log(
                    ValueError(
                        f"Invalid `norm_type={norm_type}`. Must be on of {NORMS}."
                    ),
                    logger=logger,
                )
            if norm_type == "TimeBatchNorm2d":
                norm_type = TimeBatchNorm2d
            else:
                norm_type = getattr(layer_norm_variants, norm_type)
        else:
            norm_type = norm_type

        mixer_params = {
            "ff_size": ff_size,
            "activation": activation,
            "dropout": dropout,
            "norm_type": norm_type,
            "normalize_before": normalize_before,
        }

        self.fc_hist = nn.Linear(self.input_chunk_length, self.output_chunk_length)
        self.feature_mixing_hist = _FeatureMixing(
            sequence_length=self.output_chunk_length,
            input_dim=input_dim + past_cov_dim + future_cov_dim,
            output_dim=hidden_size,
            **mixer_params,
        )
        if future_cov_dim:
            self.feature_mixing_future = _FeatureMixing(
                sequence_length=self.output_chunk_length,
                input_dim=future_cov_dim,
                output_dim=hidden_size,
                **mixer_params,
            )
        else:
            self.feature_mixing_future = None
        self.conditional_mixer = self._build_mixer(
            prediction_length=self.output_chunk_length,
            num_blocks=num_blocks,
            hidden_size=hidden_size,
            future_cov_dim=future_cov_dim,
            static_cov_dim=static_cov_dim,
            **mixer_params,
        )
        self.fc_out = nn.Linear(hidden_size, output_dim * nr_params)

        Initializes the TSMixer module for use within a Darts forecasting model.

        Parameters
        ----------
        input_dim
            Number of input target features.
        output_dim
            Number of output target features.
        past_cov_dim
            Number of past covariate features.
        future_cov_dim
            Number of future covariate features.
        static_cov_dim
            Number of static covariate features (number of target features
            (or 1 if global static covariates) * number of static covariate features).
        nr_params
            The number of parameters of the likelihood (or 1 if no likelihood is used).
        hidden_size
           Hidden state size of the TSMixer.
        ff_size
            Dimension of the feedforward network internal to the module.
        num_blocks
            Number of mixer blocks.
        activation
            Activation function to use.
        dropout
            Dropout rate for regularization.
        norm_type
            Type of normalization to use.
        normalize_before
            Whether to apply normalization before or after mixing.
        
darts/models/forecasting/tsmixer_model.py