def forward(
        self,
        inputs,
        attention_mask: Optional[torch.Tensor] = None,
        prev_state: Optional[torch.Tensor] = None,
        use_cache: bool = False,
    ) -> torch.Tensor:
        """
        Mega's exponential moving average (EMA) sub-layer applied prior to single-headed (traditional) self-attention

        Args:
            inputs (`torch.Tensor` of shape `(sequence_length, batch_size, hidden_size)`):
                Hidden state / embedding input to update via EMA based on FFT convolution
            attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Indicates which inputs are to be ignored (mostly due to padding), where elements are either 1 for *not
                masked* or 0 for *masked*
            prev_state (`torch.Tensor` of shape `(batch_size, config.ndim)`, *optional*):
                The hidden state returned from the previous timestep during incremental decoding.
            use_cache (`bool`, default `False`):
                Whether to perfom incremental decoding; uses `prev_state` as the prior timestep, and returns the
                updated EMA hidden state for use in the next step

        Returns:
            `tuple(torch.FloatTensor)` containing various elements depending on configuration ([`MegaConfig`]) and
            inputs:
            - **hidden_states** (`torch.FloatTensor` of shape `(sequence_length, batch_size, hidden_size)`) -- Hidden
              states updated by EMA, with same shapes as inputs
            - **updated_state** (*optional*, returned when `use_cache=True`) `torch.FloatTensor of shape `(batch_size,
              config.ndim)` -- The incremental EMA state for use in the next step of incremental decoding
        """

        seq_len, bsz, embed_dim = inputs.size()
        if embed_dim != self.embed_dim:
            raise ValueError(
                f"Unexpected embedding dimension received: input is {embed_dim}, model expects {self.embed_dim}"
            )

        # sequence_length X batch_size X hidden_size
        residual = inputs * self.residual_weight

        # (sequence_length x batch_size x hidden_size) -> (batch_size x hidden_size x sequence_length)
        inputs = inputs.permute(1, 2, 0)
        # mask the input: output is a tensor with 0 in the masked positions
        if attention_mask is not None:
            inputs = inputs * (attention_mask.unsqueeze(1).type_as(inputs))

        if self.bidirectional and use_cache:
            raise RuntimeError("Bidirectional EMA does not support incremental state")

        if use_cache:
            out, updated_state = self.ema_step(inputs, seq_len, past_state=prev_state)

            # (batch_size X hidden_size) -> (1 x batch_size x hidden_size)
            out = F.silu(out + residual)

            # if incremental decoding, return the new state along with the output
            return out, updated_state
        else:
            # (hidden_size x sequence_length)
            kernel = self.get_ema_kernel(seq_len)
            fft_len = seq_len
            s_index = 0
            kernel_size = kernel.size(1)
            if self.bidirectional:
                # split the kernel for each direction of EMA
                k1, k2 = torch.split(kernel, [self.embed_dim, self.embed_dim], dim=0)
                # (hidden_size X 2*sequence_length - 1)
                kernel = F.pad(k1, (kernel_size - 1, 0)) + F.pad(k2.flip(-1), (0, kernel_size - 1))
                inputs = F.pad(inputs, (kernel_size - 1, 0))
                fft_len = fft_len + kernel_size - 1
                s_index = 2 * kernel_size - 2

            ema_output = self.fft_convolution(inputs, kernel, length=fft_len)[..., s_index : s_index + seq_len]
            ema_output = ema_output.type_as(inputs)
            # (batch_size X hidden_size X sequence_length) -> (sequence_length X batch_size X hidden_size)
            gated_ema_output = F.silu(ema_output.permute(2, 0, 1) + residual)

            return gated_ema_output, None

        Mega's exponential moving average (EMA) sub-layer applied prior to single-headed (traditional) self-attention

        Args:
            inputs (`torch.Tensor` of shape `(sequence_length, batch_size, hidden_size)`):
                Hidden state / embedding input to update via EMA based on FFT convolution
            attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Indicates which inputs are to be ignored (mostly due to padding), where elements are either 1 for *not
                masked* or 0 for *masked*
            prev_state (`torch.Tensor` of shape `(batch_size, config.ndim)`, *optional*):
                The hidden state returned from the previous timestep during incremental decoding.
            use_cache (`bool`, default `False`):
                Whether to perfom incremental decoding; uses `prev_state` as the prior timestep, and returns the
                updated EMA hidden state for use in the next step

        Returns:
            `tuple(torch.FloatTensor)` containing various elements depending on configuration ([`MegaConfig`]) and
            inputs:
            - **hidden_states** (`torch.FloatTensor` of shape `(sequence_length, batch_size, hidden_size)`) -- Hidden
              states updated by EMA, with same shapes as inputs
            - **updated_state** (*optional*, returned when `use_cache=True`) `torch.FloatTensor of shape `(batch_size,
              config.ndim)` -- The incremental EMA state for use in the next step of incremental decoding
        
dpkd/transformers/src/transformers/models/mega/modeling_mega.py