docstr: Parameters:
      x: torch.tensor(bs, seq_length, dim)
      attn_mask: torch.tensor(bs, seq_length)
  
  Returns:
      sa_weights: torch.tensor(bs, n_heads, seq_length, seq_length) The attention weights ffn_output:
      torch.tensor(bs, seq_length, dim) The output of the transformer block contextualization.
docstr_tok: ["Parameters",":","x",":","torch.tensor","(","bs",",","seq_length",",","dim",")","attn_mask",":","torch.tensor","(","bs",",","seq_length",")","Returns",":","sa_weights",":","torch.tensor","(","bs",",","n_heads",",","seq_length",",","seq_length",")","The","attention","weights","ffn_output",":","torch.tensor","(","bs",",","seq_length",",","dim",")","The","output","of","the","transformer","block","contextualization","."]
code: def forward(
      self,
      x: torch.Tensor,
      attn_mask: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.Tensor] = None,
      output_attentions: bool = False,
  ) -> Tuple[torch.Tensor, ...]:
      """
      Parameters:
          x: torch.tensor(bs, seq_length, dim)
          attn_mask: torch.tensor(bs, seq_length)
  
      Returns:
          sa_weights: torch.tensor(bs, n_heads, seq_length, seq_length) The attention weights ffn_output:
          torch.tensor(bs, seq_length, dim) The output of the transformer block contextualization.
      """
      adjust_tensors_for_parallel_(x, attn_mask)
      attn_mask = prefix_attention_mask(attn_mask, dim=1, prefix_value=1)  # type: ignore
  
      # Self-Attention
      sa_output = self.attention(
          query=x,
          key=x,
          value=x,
          mask=attn_mask,
          head_mask=head_mask,
          output_attentions=output_attentions,
      )
      if output_attentions:
          (
              sa_output,
              sa_weights,
          ) = sa_output  # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
      else:  # To handle these `output_attentions` or `output_hidden_states` cases returning tuples
          if type(sa_output) is not tuple:
              raise TypeError(
                  f"sa_output must be a tuple but it is {type(sa_output)} type"
              )
  
          sa_output = sa_output[0]
      sa_output = self.attention_adapters(
          sa_output, x, self.sa_layer_norm
      )  # (bs, seq_length, dim)
  
      # Feed Forward Network
      ffn_output = self.ffn(sa_output)  # (bs, seq_length, dim)
      ffn_output: torch.Tensor = self.output_adapters(
          ffn_output, sa_output, self.output_layer_norm
      )  # (bs, seq_length, dim)
  
      output = (ffn_output,)
      if output_attentions:
          output = (sa_weights,) + output
      return output
  
file_path: adapters/src/adapters/models/distilbert/modeling_distilbert.py
imports: import math
  from typing import Optional, Tuple
  import torch
  from torch import nn
  from transformers.models.distilbert.modeling_distilbert import MultiHeadSelfAttention, TransformerBlock
  from ...composition import adjust_tensors_for_parallel, adjust_tensors_for_parallel_, match_attn_matrices_for_parallel
  from ...utils import prefix_attention_mask
  from .mixin_distilbert import DistilBertMultiHeadSelfAttentionMixin, DistilBertTransfomerBlockAdaptersMixin
label_desc: Use image node_id 1 for calling the TransformerBlockWithAdapters obj's underlying member method code with example usage: obj.forward(x, attn_mask, head_mask, output_attentions) and returns: output