def forward(
        self,
        hidden_states: torch.Tensor,
        global_hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        output_attentions: bool = False,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape *(seq_len, batch, embed_dim)*
            global_hidden_states (`torch.FloatTensor`): global token hidden states
                *(seq_len, num_global_tokens, embed_dim)*
            attention_mask (`torch.FloatTensor`): attention mask of size
                *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
        """
        residual = hidden_states
        global_residual = global_hidden_states

        hidden_states = self.self_attn_layer_norm(hidden_states)
        global_hidden_states = self.global_self_attn_layer_norm(global_hidden_states)

        if self.stagger_blocks_this_layer:
            # Pad the blocks to simulate staggering
            hidden_states, attention_mask = self.pad_local_tokens(
                hidden_states=hidden_states, attention_mask=attention_mask, block_size=self.block_size
            )

        hidden_states, global_hidden_states, attn_weights = self.self_attn(
            token_hidden_states=hidden_states,
            global_hidden_states=global_hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
        )

        if self.stagger_blocks_this_layer:
            # Undo the padding
            hidden_states = self.unpad_local_tokens(padded_hidden_states=hidden_states, block_size=self.block_size)

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        global_hidden_states = nn.functional.dropout(global_hidden_states, p=self.dropout, training=self.training)
        global_hidden_states = global_residual + global_hidden_states

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        global_residual = global_hidden_states
        global_hidden_states = self.final_layer_norm(global_hidden_states)
        global_hidden_states = self.activation_fn(self.fc1(global_hidden_states))
        global_hidden_states = nn.functional.dropout(
            global_hidden_states, p=self.activation_dropout, training=self.training
        )
        global_hidden_states = self.fc2(global_hidden_states)
        global_hidden_states = nn.functional.dropout(global_hidden_states, p=self.dropout, training=self.training)
        global_hidden_states = global_residual + global_hidden_states
        outputs = (hidden_states, global_hidden_states)

        if output_attentions:
            outputs += (attn_weights,)

        return outputs

        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape *(seq_len, batch, embed_dim)*
            global_hidden_states (`torch.FloatTensor`): global token hidden states
                *(seq_len, num_global_tokens, embed_dim)*
            attention_mask (`torch.FloatTensor`): attention mask of size
                *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
        
dpkd/transformers/src/transformers/models/pegasus_x/modeling_pegasus_x.py