def __init__(
        self,
        sequence_length: int,
        input_dim: int,
        output_dim: int,
        ff_size: int,
        activation: Callable[[torch.Tensor], torch.Tensor],
        dropout: float,
        normalize_before: bool,
        norm_type: nn.Module,
    ) -> None:
        """A module for feature mixing with flexibility in normalization and activation based on the
        `PyTorch implementation of TSMixer <https://github.com/ditschuk/pytorch-tsmixer>`_.

        This module provides options for batch normalization before or after mixing
        features, uses dropout for regularization, and allows for different activation
        functions.

        Parameters
        ----------
        sequence_length
            The length of the input sequences.
        input_dim
            The number of input channels to the module.
        output_dim
            The number of output channels from the module.
        ff_size
            The dimension of the feed-forward network internal to the module.
        activation
            The activation function used within the feed-forward network.
        dropout
            The dropout probability used for regularization.
        normalize_before
            A boolean indicating whether to apply normalization before
            the rest of the operations.
        norm_type
            The type of normalization to use.
        """
        super().__init__()

        self.projection = (
            nn.Linear(input_dim, output_dim)
            if input_dim != output_dim
            else nn.Identity()
        )
        self.norm_before = (
            norm_type((sequence_length, input_dim))
            if normalize_before
            else nn.Identity()
        )
        self.fc1 = nn.Linear(input_dim, ff_size)
        self.activation = activation
        self.dropout1 = MonteCarloDropout(dropout)
        self.fc2 = nn.Linear(ff_size, output_dim)
        self.dropout2 = MonteCarloDropout(dropout)
        self.norm_after = (
            norm_type((sequence_length, output_dim))
            if not normalize_before
            else nn.Identity()
        )
A module for feature mixing with flexibility in normalization and activation based on the
        `PyTorch implementation of TSMixer <https://github.com/ditschuk/pytorch-tsmixer>`_.

        This module provides options for batch normalization before or after mixing
        features, uses dropout for regularization, and allows for different activation
        functions.

        Parameters
        ----------
        sequence_length
            The length of the input sequences.
        input_dim
            The number of input channels to the module.
        output_dim
            The number of output channels from the module.
        ff_size
            The dimension of the feed-forward network internal to the module.
        activation
            The activation function used within the feed-forward network.
        dropout
            The dropout probability used for regularization.
        normalize_before
            A boolean indicating whether to apply normalization before
            the rest of the operations.
        norm_type
            The type of normalization to use.
        
darts/models/forecasting/tsmixer_model.py