def forward(
        self,
        input,
        padding_mask: Optional[torch.Tensor] = None,
        causal_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[torch.Tensor]] = None,
        output_attentions=False,
        use_cache=False,
    ):
        """
        Mega's self-attention block, which combines multi-headed EMA with traditional self-attention

        Args:
            input (`torch.Tensor` of shape `(sequence_length, batch_size, hidden_size)`):
                Hidden states to be updated by Mega's self-attention
            padding_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Indicates which inputs are to be ignored due to padding, where elements are either 1 for *not masked*
                or 0 for *masked*
            causal_mask (`torch.LongTensor` of shape `(sequence_length, sequence_length)`, *optional*):
                Indicates which inputs are to be ignored due to causal attention, where elements are either 1 for *not
                masked* or 0 for *masked*
            past_key_values (`tuple(torch.Tensor)`, *optional*):
                The hidden states returned from the previous timestep during incremental decoding; expects that
                self-attention key, value, and EMA states are the first 3 entries in the tuple
            output_attentions (`bool`, default `False`):
                Whether to return self-attention weights
            use_cache (`bool`, default `False`):
                Whether to perfom incremental decoding; uses `past_key_values` as prior state, and returns the updated
                states 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 from target sequence updated by Mega's self-attention
            - **attn_weights** (*optional*, returned when `output_attentions=True`) `torch.FloatTensor` of shape
              `(batch_size, 1, sequence_length, sequence_length)` -- The self-attention weights corresponding to how
              each token in the input sequence attends to every other token
            - **self_key** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              sequence_length, config.shared_representation_size)` -- The self-attention key state for use in the next
              step of incremental decoding
            - **self_value** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              sequence_length, config.hidden_size)` -- The self-attention value state for use in the next step of
              incremental decoding
            - **self_ema_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 = input.size()
        if embed_dim != self.config.hidden_size:
            raise ValueError(f"Input embedding dimension should be {self.config.hidden_size}; received {embed_dim}")

        # store inputs for residual connection and handle pre-norm if requested
        residual = input
        if self.config.normalize_before_mega:
            input = self.norm(input)

        # (sequence_length X batch_size X hidden_size) -> (sequence_length X batch_size X intermediate_size)
        value = self.activation(self.v_proj(input))

        # unpack the incremental state if provided
        # assumed to be (self K, self V, self EMA state, cross K, cross V)
        # also assumes that incremental decoding is working one token at a time, so input sequence length must be 1
        if self.config.is_decoder and (past_key_values is not None):
            if seq_len > 1:
                raise ValueError(f"Incremental decoding only supports self sequence length of 1; received {seq_len}")
            # the first 3 items in the saved states will be these regardless of whether cross-attention is present
            prev_self_key, prev_self_value, prev_ema_state = past_key_values[0:3]
        else:
            prev_self_key = prev_self_value = prev_ema_state = None

        # ema output is (sequence_length x batch_size x hidden_size)
        # updated_ema_state will be None if use_cache=False; otherwise (batch_size, config.ndim)
        ema_out, updated_ema_state = self.ema_gate(
            input, attention_mask=padding_mask, prev_state=prev_ema_state, use_cache=use_cache
        )
        ema_out = self.dropout(ema_out)

        # (sequence_length X batch_size X hidden_size)
        # -> (sequence_length X batch_size X 2*hidden_size + config.shared_representation_size + config.intermediate_size)
        # - residual_weight -> sigmoid -> applied to residual connection in torch.addcmul
        # - query_key_gates -> split into two components: query_key becomes query and key for attention input, gates becomes gating for self-attention output
        # - intermediate_state -> added to weighted attention output, sent through activation, and has inputs subtracted during
        #   torch.addcmul to create the final layer output
        base = self.mx_proj(ema_out)
        residual_weight, query_key_gates, intermediate_state = torch.split(
            base,
            [
                self.config.hidden_size,
                self.config.shared_representation_size + self.config.intermediate_size,
                self.config.hidden_size,
            ],
            dim=-1,
        )

        # (sequence_length X batch_size X hidden_size)
        residual_weight = torch.sigmoid(residual_weight)

        # (sequence_length X batch_size X shared_representation_size + intermediate_size)
        query_key_gates = F.silu(query_key_gates)

        # split into two different tensors: one for Q/K usage and the other for gating self-attention
        query_key, attention_gate = torch.split(
            query_key_gates, [self.config.shared_representation_size, self.config.intermediate_size], dim=-1
        )

        # (sequence_length X batch_size X shared_representation_size)
        # -> (sequence_length X batch_size X 1 X shared_representation_size)
        # -> (sequence_length X batch_size X 2 X shared_representation_size)
        query_key = query_key.unsqueeze(2) * self.qk_weight + self.qk_bias

        # (sequence_length X batch_size X 2 X shared_representation_size)
        # -> 2 tensors of (sequence_length X batch_size X shared_representation_size)
        query, key = torch.unbind(query_key, dim=2)

        # (sequence_length X batch_size X dimension)
        # -> (batch_size X sequence_length X dimension)
        # where `dimension` is either shared_representation_size (queries and keys) or intermediate_size (values)
        query = query.transpose(0, 1)
        key = key.transpose(0, 1)
        value = value.transpose(0, 1)

        if self.config.is_decoder:
            # combine history and current to save updated state (if history is provided)
            # when chunking is applied, the past states will be None at the end of the chunk, in
            # which case, proceed as if no K/V history had been provided
            # saved states are stored with shape (batch_size X sequence_length X dimension)
            if prev_self_key is not None:
                key = torch.cat([prev_self_key, key], dim=1)
            if prev_self_value is not None:
                value = torch.cat([prev_self_value, value], dim=1)

            # if not chunking, store as-is
            if not self.config.use_chunking:
                updated_self_key = key
                updated_self_value = value
            else:
                curr_len = key.size(1) % self.config.chunk_size
                if curr_len == 0:
                    # if we're chunking and have reached the end of a chunk, wipe out the saved state
                    updated_self_key = None
                    updated_self_value = None
                else:
                    updated_self_key = key
                    updated_self_value = value

        ctx_len = key.size(1)  # potentially differs from seq_len because of incremental decoding
        if not self.config.use_chunking:
            # if we're not chunking, treat the entire sequence as one long chunk
            # (batch_size X sequence_length X dimension) -> (batch_size X 1 X sequence_length X dimension)
            query = query.unsqueeze(1)
            key = key.unsqueeze(1)
            value = value.unsqueeze(1)
            if padding_mask is not None:
                # (batch_size X sequence_length) -> (batch_size X 1 X sequence_length)
                padding_mask = padding_mask.unsqueeze(1)
        else:
            # otherwise, split the sequences in the batch into `n_chunks` chunks of size `chunk_size`
            if seq_len < self.config.chunk_size:
                query = query.unsqueeze(1)
            else:
                # (batch_size X sequence_length X dimension) -> (batch_size X n_chunks X chunk_size X dimension)
                n_chunks = seq_len // self.config.chunk_size
                query = query.reshape(bsz, n_chunks, self.config.chunk_size, self.config.shared_representation_size)

            if ctx_len < self.config.chunk_size:
                key = key.unsqueeze(1)
                value = value.unsqueeze(1)
                if padding_mask is not None:
                    padding_mask = padding_mask.unsqueeze(1)
            else:
                # (batch_size X sequence_length X dimension) -> (batch_size X n_chunks X chunk_size X dimension)
                n_chunks = ctx_len // self.config.chunk_size
                key = key.reshape(bsz, n_chunks, self.config.chunk_size, self.config.shared_representation_size)
                value = value.reshape(bsz, n_chunks, self.config.chunk_size, self.config.intermediate_size)
                if padding_mask is not None:
                    padding_mask = padding_mask.view(bsz, n_chunks, self.config.chunk_size)

        # this is in the original Mega implementation to work around fork/join parallelism not supporting optional types
        if padding_mask is not None and padding_mask.dim() == 0:
            padding_mask = None

        attn_weights = self.attention_function(query, key, padding_mask=padding_mask, causal_mask=causal_mask)

        value = self.hidden_dropout(value, batch_first=True)
        kernel = self.attention_dropout(attn_weights)

        # (batch_size x n_chunks x chunk_size x intermediate_size) -> (sequence_length X batch_size X intermediate_size)
        weighted_self_output = (
            torch.matmul(kernel, value).view(bsz, seq_len, self.config.intermediate_size).transpose(0, 1)
        )

        # (sequence_length X batch_size X intermediate_size) -> (sequence_length X batch_size X hidden_size)
        weighted_self_output = self.activation(intermediate_state + self.h_proj(weighted_self_output * attention_gate))
        weighted_self_output = self.dropout(weighted_self_output)
        # (sequence_length X batch_size X hidden_size)
        out = torch.addcmul(residual, residual_weight, weighted_self_output - residual)

        if not self.config.normalize_before_mega:
            out = self.norm(out)

        return_values = (out, attn_weights) if output_attentions else (out,)

        if self.config.is_decoder:
            return_values = return_values + (updated_self_key, updated_self_value, updated_ema_state)

        return return_values

        Mega's self-attention block, which combines multi-headed EMA with traditional self-attention

        Args:
            input (`torch.Tensor` of shape `(sequence_length, batch_size, hidden_size)`):
                Hidden states to be updated by Mega's self-attention
            padding_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Indicates which inputs are to be ignored due to padding, where elements are either 1 for *not masked*
                or 0 for *masked*
            causal_mask (`torch.LongTensor` of shape `(sequence_length, sequence_length)`, *optional*):
                Indicates which inputs are to be ignored due to causal attention, where elements are either 1 for *not
                masked* or 0 for *masked*
            past_key_values (`tuple(torch.Tensor)`, *optional*):
                The hidden states returned from the previous timestep during incremental decoding; expects that
                self-attention key, value, and EMA states are the first 3 entries in the tuple
            output_attentions (`bool`, default `False`):
                Whether to return self-attention weights
            use_cache (`bool`, default `False`):
                Whether to perfom incremental decoding; uses `past_key_values` as prior state, and returns the updated
                states 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 from target sequence updated by Mega's self-attention
            - **attn_weights** (*optional*, returned when `output_attentions=True`) `torch.FloatTensor` of shape
              `(batch_size, 1, sequence_length, sequence_length)` -- The self-attention weights corresponding to how
              each token in the input sequence attends to every other token
            - **self_key** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              sequence_length, config.shared_representation_size)` -- The self-attention key state for use in the next
              step of incremental decoding
            - **self_value** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              sequence_length, config.hidden_size)` -- The self-attention value state for use in the next step of
              incremental decoding
            - **self_ema_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