def forward(
        self,
        query,
        key: Optional[torch.Tensor],
        value: Optional[torch.Tensor],
        key_padding_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Gated cross-attention used in Mega

        Args:
            query (`torch.Tensor` of shape `(target_sequence_length, batch_size, hidden_size)`):
                The self (or target) sequence input used as query inputs for cross-attention
            key (`torch.Tensor` of shape `(source_sequence_length, batch_size, hidden_size)`):
                The cross (or source) sequence input with shape used as keys in cross-attention
            value (`torch.Tensor` of shape `(source_sequence_length, batch_size, hidden_size)`):
                The cross (or source) sequence input with shape used as values in cross-attention
            key_padding_mask (`torch.LongTensor` of shape `(batch_size, source_sequence_length)`, *optional*):
                Padding mask corresponding to the source sequence, where entries are 1 for *not masked* and 0 for
                *masked* tokens
            past_key_values (`tuple(torch.FloatTensor)`, *optional*):
                If provided, the hidden state returned from the previous timestep during incremental decoding; expects
                that prior cross-attention keys and values will be the last two items in the tuple
            output_attentions (`bool`, defaults to `False`):
                Whether or not to return the cross-attention weights.
            use_cache (`bool`, defaults to `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 `(target_sequence_length, batch_size, hidden_size)`) --
              Hidden states from target sequence updated by gated cross-attention
            - **attn_weights** (*optional*, returned when `output_attentions=True`) `torch.FloatTensor` of shape
              `(batch_size, source_sequence_length, target_sequence_length)` -- The pairwise cross-attention weights
              corresponding to each token in the source and target sequences
            - **cross_key** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              source_sequence_length, config.shared_representation_size)` -- The cross-attention key state for use in
              the next step of incremental decoding
            - **cross_value** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              source_sequence_length, config.hidden_size)` -- The cross-attention value state for use in the next step
              of incremental decoding
        """

        seq_len, bsz, embed_dim = query.size()
        if embed_dim != self.config.hidden_size:
            raise ValueError(
                f"Unexpected embedding dimension received: input is {embed_dim} but expected {self.config.hidden_size}"
            )

        if past_key_values is not None:
            # make sure the inputs only have a sequence length of 1 if we're doing incremental decoding
            if seq_len != 1:
                raise ValueError(f"Incremental decoding requested with self-sequence length > 1: {seq_len}")
            # expect past_key_values to have (self_key, self_value, self_ema, cross_key, cross_value)
            prev_cross_key, prev_cross_value = past_key_values[-2:]
            key = value = None

            # use the self-attention cache to get the position id of the current step
            prev_self_key = past_key_values[0]
            num_incremental_steps = prev_self_key.size(1) + 1
        else:
            prev_cross_key = prev_cross_value = None
            # we still need the position id if we're doing incremental decoding (past_key_values will be None for the first step)
            num_incremental_steps = 0 if use_cache and (seq_len == 1) else None

        full_query = query
        if self.prenorm:
            full_query = self.norm(full_query)

        # (target_sequence_length X batch_size X 2*hidden_size + shared_representation_size)
        query_projected = self.q_proj(full_query)
        # split the query projections into separate components
        # - residual_weight is passed through sigmoid and sent through elementwise multiplication to the gated/weighted targets prior to being added to the query directly
        # - target_gate is a silu-gated tensor that is multiplied by the attention-weighted target below prior to residual connection
        # - attention_query is the part that is passed to the attention function
        residual_weight, target_gate, attention_query = torch.split(
            query_projected,
            [self.config.hidden_size, self.config.hidden_size, self.config.shared_representation_size],
            dim=-1,
        )

        # (target_sequence_length X batch_size X hidden_size)
        residual_weight = torch.sigmoid(residual_weight)
        target_gate = F.silu(target_gate)

        if key is None:
            if value is not None:
                raise ValueError("Key and value must be `None` simultaneously")
            projected_key = projected_value = None
        else:
            # (source_sequence_length X batch_size X shared_representation_size)
            projected_key = self.k_proj(key)
            # (source_sequence_length X batch_size X hidden_size)
            projected_value = self.activation(self.v_proj(key))

        # (target_sequence_length X batch_size X shared_representation_size)
        # -> (batch_size X target_sequence_length X shared_representation_size)
        attention_query = attention_query.transpose(0, 1)
        if projected_key is not None:
            projected_key = projected_key.transpose(0, 1)
        if projected_value is not None:
            projected_value = projected_value.transpose(0, 1)

        # if we're doing incremental decoding, k and v are None and need to be overwritten with past values
        if past_key_values is not None:
            projected_key = prev_cross_key
            projected_value = prev_cross_value

        # if we're returning the cache for later use, store these now for later return (can be done without having past_key_values provided)
        if use_cache:
            updated_cross_key = projected_key
            updated_cross_value = projected_value

        ctx_len = projected_key.size(1)
        # This is part of a workaround to get around fork/join parallelism
        # not supporting Optional types.
        if key_padding_mask is not None and key_padding_mask.dim() == 0:
            key_padding_mask = None

        if key_padding_mask is not None:
            if key_padding_mask.size(0) != bsz:
                raise ValueError("Key padding mask does not align on the batch dimension")
            if key_padding_mask.size(1) != ctx_len:
                raise ValueError("Key padding mask does not align on the sequence length dimension")

        if self.attention_activation == "softmax":
            attn_weights = self.softmax_attention(
                attention_query, projected_key, key_padding_mask, num_incremental_steps
            )
        else:
            attn_weights = self.element_attention(
                attention_query, projected_key, key_padding_mask, num_incremental_steps
            )

        projected_value = self.hidden_dropout(projected_value, batch_first=True)
        kernel = self.attention_dropout(attn_weights)
        # (batch_size X target_sequence_length X hidden_size)
        # -> (target_sequence_length X batch_size X hidden_size)
        weighted_targets = torch.bmm(kernel, projected_value).transpose(0, 1)
        # (target_sequence_length X batch_size X hidden_size)
        weighted_targets = self.activation(self.h_proj(weighted_targets * target_gate))
        weighted_targets = self.dropout(weighted_targets)
        out = torch.addcmul(query, residual_weight, weighted_targets - query)

        if not self.prenorm:
            out = self.norm(out)

        outputs = (out, attn_weights) if output_attentions else (out,)
        if use_cache:
            outputs = outputs + (updated_cross_key, updated_cross_value)

        return outputs

        Gated cross-attention used in Mega

        Args:
            query (`torch.Tensor` of shape `(target_sequence_length, batch_size, hidden_size)`):
                The self (or target) sequence input used as query inputs for cross-attention
            key (`torch.Tensor` of shape `(source_sequence_length, batch_size, hidden_size)`):
                The cross (or source) sequence input with shape used as keys in cross-attention
            value (`torch.Tensor` of shape `(source_sequence_length, batch_size, hidden_size)`):
                The cross (or source) sequence input with shape used as values in cross-attention
            key_padding_mask (`torch.LongTensor` of shape `(batch_size, source_sequence_length)`, *optional*):
                Padding mask corresponding to the source sequence, where entries are 1 for *not masked* and 0 for
                *masked* tokens
            past_key_values (`tuple(torch.FloatTensor)`, *optional*):
                If provided, the hidden state returned from the previous timestep during incremental decoding; expects
                that prior cross-attention keys and values will be the last two items in the tuple
            output_attentions (`bool`, defaults to `False`):
                Whether or not to return the cross-attention weights.
            use_cache (`bool`, defaults to `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 `(target_sequence_length, batch_size, hidden_size)`) --
              Hidden states from target sequence updated by gated cross-attention
            - **attn_weights** (*optional*, returned when `output_attentions=True`) `torch.FloatTensor` of shape
              `(batch_size, source_sequence_length, target_sequence_length)` -- The pairwise cross-attention weights
              corresponding to each token in the source and target sequences
            - **cross_key** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              source_sequence_length, config.shared_representation_size)` -- The cross-attention key state for use in
              the next step of incremental decoding
            - **cross_value** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
              source_sequence_length, config.hidden_size)` -- The cross-attention value state for use in the next step
              of incremental decoding
        
dpkd/transformers/src/transformers/models/mega/modeling_mega.py