class_docstr: Multi-headed attention from 'Attention Is All You Need' paper
class_docstr_tok: ["Multi-headed","attention","from","'Attention","Is","All","You","Need","'","paper"]
docstr: Input shape: Batch x Time x Channel
docstr_tok: ["Input","shape",":","Batch","x","Time","x","Channel"]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      key_value_states: Optional[torch.Tensor] = None,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      attention_mask: Optional[torch.Tensor] = None,
      layer_head_mask: Optional[torch.Tensor] = None,
      output_attentions: bool = False,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      """Input shape: Batch x Time x Channel"""
  
      # if key_value_states are provided this layer is used as a cross-attention layer
      # for the decoder
      is_cross_attention = key_value_states is not None
  
      bsz, tgt_len, _ = hidden_states.size()
  
      # get query proj
      query_states = self.q_proj(hidden_states) * self.scaling
      # get key, value proj
      # `past_key_value[0].shape[2] == key_value_states.shape[1]`
      # is checking that the `sequence_length` of the `past_key_value` is the same as
      # the provided `key_value_states` to support prefix tuning
      if (
          is_cross_attention
          and past_key_value is not None
          and past_key_value[0].shape[2] == key_value_states.shape[1]
      ):
          # reuse k,v, cross_attentions
          key_states = past_key_value[0]
          value_states = past_key_value[1]
      elif is_cross_attention:
          # cross_attentions
          key_states = self._shape(
              self.k_proj(key_value_states), -1, bsz
          )
          value_states = self._shape(
              self.v_proj(key_value_states), -1, bsz
          )
      elif past_key_value is not None:
          # reuse k, v, self_attention
          key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
          value_states = self._shape(
              self.v_proj(hidden_states), -1, bsz
          )
          key_states = torch.cat([past_key_value[0], key_states], dim=2)
          value_states = torch.cat(
              [past_key_value[1], value_states], dim=2
          )
      else:
          # self_attention
          key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
          value_states = self._shape(
              self.v_proj(hidden_states), -1, bsz
          )
  
      if self.is_decoder:
          # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
          # Further calls to cross_attention layer can then reuse all cross-attention
          # key/value_states (first "if" case)
          # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
          # all previous decoder key/value_states. Further calls to uni-directional self-attention
          # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
          # if encoder bi-directional self-attention `past_key_value` is always `None`
          past_key_value = (key_states, value_states)
  
      proj_shape = (bsz * self.num_heads, -1, self.head_dim)
      query_states = self._shape(query_states, tgt_len, bsz).view(
          *proj_shape
      )
      key_states = key_states.reshape(*proj_shape)
      value_states = value_states.reshape(*proj_shape)
  
      src_len = key_states.size(1)
      attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
  
      if attn_weights.size() != (
          bsz * self.num_heads,
          tgt_len,
          src_len,
      ):
          raise ValueError(
              f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
              f" {attn_weights.size()}"
          )
  
      if attention_mask is not None:
          if attention_mask.size() != (bsz, 1, tgt_len, src_len):
              raise ValueError(
                  f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
              )
          attn_weights = (
              attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              + attention_mask
          )
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      attn_weights = nn.functional.softmax(attn_weights, dim=-1)
  
      if layer_head_mask is not None:
          if layer_head_mask.size() != (self.num_heads,):
              raise ValueError(
                  f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
                  f" {layer_head_mask.size()}"
              )
          attn_weights = layer_head_mask.view(
              1, -1, 1, 1
          ) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      if output_attentions:
          # this operation is a bit awkward, but it's required to
          # make sure that attn_weights keeps its gradient.
          # In order to do so, attn_weights have to be reshaped
          # twice and have to be reused in the following
          attn_weights_reshaped = attn_weights.view(
              bsz, self.num_heads, tgt_len, src_len
          )
          attn_weights = attn_weights_reshaped.view(
              bsz * self.num_heads, tgt_len, src_len
          )
      else:
          attn_weights_reshaped = None
  
      attn_probs = nn.functional.dropout(
          attn_weights, p=self.dropout, training=self.training
      )
  
      attn_output = torch.bmm(attn_probs, value_states)
  
      if attn_output.size() != (
          bsz * self.num_heads,
          tgt_len,
          self.head_dim,
      ):
          raise ValueError(
              f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
              f" {attn_output.size()}"
          )
  
      attn_output = attn_output.view(
          bsz, self.num_heads, tgt_len, self.head_dim
      )
      attn_output = attn_output.transpose(1, 2)
  
      # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
      # partitioned across GPUs when using tensor-parallelism.
      attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
  
      attn_output = self.out_proj(attn_output)
  
      return attn_output, attn_weights_reshaped, past_key_value
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss
  from ...activations import ACT2FN
  from ...integrations.deepspeed import is_deepspeed_zero3_enabled
  from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask
  from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput, Seq2SeqModelOutput
  from ...modeling_utils import PreTrainedModel
  from ...utils import add_code_sample_docstrings, add_end_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_m2m_100 import M2M100Config
label_desc: Use image node_id 3 for calling the M2M100Attention obj's underlying member method code with example usage: obj.forward(hidden_states, key_value_states, past_key_value, attention_mask, layer_head_mask, output_attentions) and returns: attn_output,  attn_weights_reshaped,  past_key_value