class_docstr: Implements MultiHeadAttention; this is the core workhorse of the Transformer.
  
  See Vaswani (2017) for an extensive description.
class_docstr_tok: ["Implements","MultiHeadAttention",";","this","is","the","core","workhorse","of","the","Transformer",".","See","Vaswani","(","2017",")","for","an","extensive","description","."]
docstr: Forward pass.
  
  :param query: attention query
  :param key: attention key
  :param value: attention value
  :param mask: tensor in which True means that we are allowing attention and False
    means we are blocking it. Mask is:
    - [B, key_len] (encoder self-attn and decoder enc/dec attn)
    - [B, query_len, key_len] (decoder self-attn)
    - [B, 1, key_len] (decoder self-attn with incr_state caching)
  :param incr_state: dictionary with values representing the previous states of
    the key, value, and mask
  :param static_kv: True if the key and value are held constant during decoding
    (as in encoder/decoder attention)
  :return: (
    final attended tensor,
    new incremental state,
    key/value-multiplied tensor before softmax,
  )
docstr_tok: ["Forward","pass",".",":","param","query",":","attention","query",":","param","key",":","attention","key",":","param","value",":","attention","value",":","param","mask",":","tensor","in","which","True","means","that","we","are","allowing","attention","and","False","means","we","are","blocking","it",".","Mask","is",":","-","[","B",",","key_len","]","(","encoder","self-attn","and","decoder","enc\/dec","attn",")","-","[","B",",","query_len",",","key_len","]","(","decoder","self-attn",")","-","[","B",",","1",",","key_len","]","(","decoder","self-attn","with","incr_state","caching",")",":","param","incr_state",":","dictionary","with","values","representing","the","previous","states","of","the","key",",","value",",","and","mask",":","param","static_kv",":","True","if","the","key","and","value","are","held","constant","during","decoding","(","as","in","encoder\/decoder","attention",")",":","return",":","(","final","attended","tensor",",","new","incremental","state",",","key\/value-multiplied","tensor","before","softmax",",",")"]
code: def forward(  # type: ignore
      # TODO: remove type ignore with pytorch 1.5:
      # https://github.com/pytorch/pytorch/pull/31057
      self,
      query: torch.Tensor,
      key: Optional[torch.Tensor] = None,
      value: Optional[torch.Tensor] = None,
      mask: torch.Tensor = None,
      incr_state: Optional[Dict[str, torch.Tensor]] = None,
      static_kv: bool = False,
      **kwargs,
  ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
      """
      Forward pass.
  
      :param query: attention query
      :param key: attention key
      :param value: attention value
      :param mask: tensor in which True means that we are allowing attention and False
        means we are blocking it. Mask is:
        - [B, key_len] (encoder self-attn and decoder enc/dec attn)
        - [B, query_len, key_len] (decoder self-attn)
        - [B, 1, key_len] (decoder self-attn with incr_state caching)
      :param incr_state: dictionary with values representing the previous states of
        the key, value, and mask
      :param static_kv: True if the key and value are held constant during decoding
        (as in encoder/decoder attention)
      :return: (
        final attended tensor,
        new incremental state,
        key/value-multiplied tensor before softmax,
      )
      """
  
      batch_size, query_len, dim = query.size()
      assert (
          dim == self.dim
      ), "Dimensions do not match: {} query vs {} configured".format(
          dim, self.dim
      )
      assert mask is not None, "Mask is None, please specify a mask"
      n_heads = self.n_heads
      dim_per_head = dim // n_heads
      scale = math.sqrt(dim_per_head)
  
      def prepare_head(tensor):
          # input is [batch_size, seq_len, n_heads * dim_per_head]
          # output is [batch_size * n_heads, seq_len, dim_per_head]
          bsz, seq_len, _ = tensor.size()
          tensor = tensor.view(
              batch_size, tensor.size(1), n_heads, dim_per_head
          )
          tensor = (
              tensor.transpose(1, 2)
              .contiguous()
              .view(batch_size * n_heads, seq_len, dim_per_head)
          )
          return tensor
  
      # q, k, v are the transformed values
      if key is None and value is None:
          # self attention
          key = value = query
          _, _key_len, dim = query.size()
      elif value is None:
          # key and value are the same, but query differs
          # self attention
          value = key
  
      assert key is not None  # let mypy know we sorted this
      _, _key_len, dim = key.size()
  
      q = prepare_head(self.q_lin(query))
  
      # Prepend incremental states. For each of the key, value, and mask, see if
      # a previous incremental state exists, and if so, reshape it to match the shape
      # of the new state. Concatenate the previous and new states to match what the
      # full state would have been if we had not cached. (If we are using static_kv,
      # these three states are unchanging, so just re-use the cached states.)
      if incr_state is None:
          incr_state = {}
      if "prev_key" in incr_state:
          prev_key = incr_state["prev_key"].view(
              batch_size * n_heads, -1, dim_per_head
          )
          if static_kv:
              k = prev_key
          else:
              k = torch.cat(
                  [prev_key, prepare_head(self.k_lin(key))], dim=1
              )
      else:
          k = prepare_head(self.k_lin(key))
      if "prev_value" in incr_state:
          prev_value = incr_state["prev_value"].view(
              batch_size * n_heads, -1, dim_per_head
          )
          if static_kv:
              v = prev_value
          else:
              v = torch.cat(
                  [prev_value, prepare_head(self.v_lin(value))], dim=1
              )
      else:
          v = prepare_head(self.v_lin(value))
      if "prev_mask" in incr_state:
          if static_kv:
              mask = incr_state["prev_mask"]
          else:
              # Mask will be of size (B x query_len x key_len)
              # During incremental decoding the query will only represent the next token,
              # whereas the key/value will represent the entire sequence thus far.
              # In such a case, we only want to look at the last element of the mask in the query dimension.
              prev_mask = incr_state["prev_mask"][:, -query_len:, :]
              mask = torch.cat([prev_mask, mask], dim=2)
              # Prepend along the key_len dimension (analogous to incr_state['prev_key'])
  
      # Save new incremental states. We reshape to allow for reordering along batch
      # dimension.
      new_incr_state = {
          "prev_key": k.view(batch_size, n_heads, -1, dim_per_head),
          "prev_value": v.view(batch_size, n_heads, -1, dim_per_head),
          "prev_mask": mask,
      }
  
      full_key_len = k.size(1)
      dot_prod = q.div_(scale).bmm(k.transpose(1, 2))
      # [B * n_heads, query_len, key_len]
      attn_mask = (
          (mask == 0)
          .view(batch_size, 1, -1, full_key_len)
          .repeat(1, n_heads, 1, 1)
          .expand(batch_size, n_heads, query_len, full_key_len)
          .view(batch_size * n_heads, query_len, full_key_len)
      )
      assert attn_mask.shape == dot_prod.shape
      dot_prod.masked_fill_(attn_mask, neginf(dot_prod.dtype))
  
      attn_weights = F.softmax(
          dot_prod, dim=-1, dtype=torch.float  # type: ignore
      ).type_as(query)
      attn_weights = self.attn_dropout(
          attn_weights
      )  # --attention-dropout
  
      attentioned = attn_weights.bmm(v)
      attentioned = (
          attentioned.type_as(query)
          .view(batch_size, n_heads, query_len, dim_per_head)
          .transpose(1, 2)
          .contiguous()
          .view(batch_size, query_len, dim)
      )
  
      out = self.out_lin(attentioned)
  
      return out, new_incr_state, dot_prod
  
file_path: ParlAI/parlai/agents/transformer/modules/attention.py
imports: import math
  from typing import Dict, Tuple, Optional, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from parlai.core.opt import Opt
  from parlai.core.params import default
  from parlai.utils.torch import neginf
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, mask, incr_state, static_kv) and returns: out,  new_incr_state,  dot_prod, tensor