src/transformers/models/pegasus_x/modeling_pegasus_x.py/PegasusXGlobalLocalAttention/forward
class PegasusXGlobalLocalAttention:    def forward(
        self,
        token_hidden_states: torch.Tensor,
        global_hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        """Input shape: Batch x Time x Channel"""
        dim = DimensionInfo(
            batch_size=token_hidden_states.shape[0],
            seq_len=token_hidden_states.shape[1],
            block_size=self.block_size,
            num_heads=self.num_heads,
            hidden_dim=token_hidden_states.shape[2],
            dim_per_head=self.head_dim,
            num_blocks=token_hidden_states.shape[1] // self.block_size,
            global_len=global_hidden_states.shape[1],
            padded_seq_len=token_hidden_states.shape[1],
        )

        # [batch_size, num_heads, padded_seq_len, dim_per_head]
        local_q = self._shape(
            self.q_proj(token_hidden_states) * self.scaling,
            seq_len=dim.padded_seq_len,
            bsz=dim.batch_size,
        )
        local_k = self._shape(
            self.k_proj(token_hidden_states),
            seq_len=dim.padded_seq_len,
            bsz=dim.batch_size,
        )
        local_v = self._shape(
            self.v_proj(token_hidden_states),
            seq_len=dim.padded_seq_len,
            bsz=dim.batch_size,
        )

        # [batch_size, num_heads, global_len, dim_per_head]
        global_q = self._shape(
            self.q_proj(global_hidden_states) * self.scaling,
            seq_len=dim.global_len,
            bsz=dim.batch_size,
        )
        global_k = self._shape(
            self.k_proj(global_hidden_states),
            seq_len=dim.global_len,
            bsz=dim.batch_size,
        )
        global_v = self._shape(
            self.v_proj(global_hidden_states),
            seq_len=dim.global_len,
            bsz=dim.batch_size,
        )

        global_attn_output, global_attn_probs = self.compute_global_attention_representations(
            global_q=global_q,
            global_k=global_k,
            global_v=global_v,
            local_k=local_k,
            local_v=local_v,
            mask=attention_mask,
            dim=dim,
        )
        local_attn_output, local_attn_probs = self.compute_local_attention_representations(
            global_k=global_k,
            global_v=global_v,
            local_q=local_q,
            local_k=local_k,
            local_v=local_v,
            mask=attention_mask,
            dim=dim,
        )

        # [batch_size, global_len, hidden_dim]
        global_attn_output = (
            global_attn_output.transpose(1, 2).contiguous().view(dim.batch_size, dim.global_len, dim.hidden_dim)
        )
        # [batch_size, global_len, hidden_dim]
        global_attn_output = self.out_proj(global_attn_output)
        # [batch_size, num_heads, block_size, num_heads, dim_per_head]
        local_attn_output = local_attn_output.permute(0, 2, 3, 1, 4).contiguous()
        # [batch_size, padded_seq_len, hidden_dim]
        local_attn_output = local_attn_output.view(dim.batch_size, dim.padded_seq_len, dim.hidden_dim)
        # [batch_size, padded_seq_len, hidden_dim]
        local_attn_output = self.out_proj(local_attn_output)

        if output_attentions:
            attn_probs = {"global": global_attn_probs, "local": local_attn_probs}
        else:
            attn_probs = None

        return local_attn_output, global_attn_output, attn_probs