def __init__(self, input_dim: int, eps=1e-5, affine=True):
        """Reversible Instance Normalization based on [1]

        Parameters
        ----------
        input_dim
            The dimension of the input axis being normalized
        eps
            The epsilon value for numerical stability
        affine
            Whether to apply an affine transformation after normalization

        References
        ----------
        .. [1] Kim et al. "Reversible Instance Normalization for Accurate Time-Series Forecasting against
                Distribution Shift" International Conference on Learning Representations (2022)
        """

        super().__init__()
        self.input_dim = input_dim
        self.eps = eps
        self.affine = affine

        if self.affine:
            self.affine_weight = nn.Parameter(torch.ones(self.input_dim))
            self.affine_bias = nn.Parameter(torch.zeros(self.input_dim))
Reversible Instance Normalization based on [1]

        Parameters
        ----------
        input_dim
            The dimension of the input axis being normalized
        eps
            The epsilon value for numerical stability
        affine
            Whether to apply an affine transformation after normalization

        References
        ----------
        .. [1] Kim et al. "Reversible Instance Normalization for Accurate Time-Series Forecasting against
                Distribution Shift" International Conference on Learning Representations (2022)
        
darts/models/components/layer_norm_variants.py