class_docstr: Refer Attention is All You Need for more details.
  
  This is a simplified implementation of multihead attention
  optimized for exporting using torchscript. Usage of nn.Linear() instead of
  F.Linear() helps to quantize the linear layers.
  
  Query represents the output from last decoder step. Key and Values are obtained from
  encoder. Attention weights are obtained from the dot product of query and key.
  Attention weights multiplied by the value gives output.
class_docstr_tok: ["Refer","Attention","is","All","You","Need","for","more","details",".","This","is","a","simplified","implementation","of","multihead","attention","optimized","for","exporting","using","torchscript",".","Usage","of","nn.Linear","(",")","instead","of","F.Linear","(",")","helps","to","quantize","the","linear","layers",".","Query","represents","the","output","from","last","decoder","step",".","Key","and","Values","are","obtained","from","encoder",".","Attention","weights","are","obtained","from","the","dot","product","of","query","and","key",".","Attention","weights","multiplied","by","the","value","gives","output","."]
code: def forward(
      self,
      query: Tensor,
      key: Tensor,
      value: Tensor,
      key_padding_mask: Optional[Tensor],
      need_weights: bool,
      incremental_state: Optional[Dict[str, Tensor]] = None,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      target_len, bsz, embed_dim = query.size()
      src_len = key.size(0)
      assert embed_dim == self.embed_dim, (
          str(embed_dim) + " != " + str(self.embed_dim)
      )
      assert key is not None
      assert value is not None
  
      if incremental_state is not None:
          prev_key = self._get_input_buffer(
              incremental_state, "prev_key"
          )
      else:
          prev_key = None
  
      bsz_X_num_heads = bsz * self.num_heads
  
      q = self.q_proj(query)
      q *= self.scaling
      q = (
          q.contiguous()
          .view(target_len, bsz_X_num_heads, self.head_dim)
          .transpose(0, 1)
      )
  
      if prev_key is not None and incremental_state is not None:
          # This happens if its incremental decoding and prev time step has been
          # cached. This condition won't be true for the first step in
          # incremental decoding.
          k = prev_key.view(bsz_X_num_heads, -1, self.head_dim)
  
          prev_value = self._get_input_buffer(
              incremental_state, "prev_value"
          )
          assert prev_value is not None
  
          v = prev_value.view(bsz_X_num_heads, -1, self.head_dim)
  
      else:
          # We will recompute key and value for all regular training and
          # for first step of incremental decoding
          k = self.k_proj(key)
  
          k = (
              k.contiguous()
              .view(-1, bsz_X_num_heads, self.head_dim)
              .transpose(0, 1)
          )
  
          v = self.v_proj(value)
          v = (
              v.contiguous()
              .view(-1, bsz_X_num_heads, self.head_dim)
              .transpose(0, 1)
          )
  
          # incremental state needs to be set only for the first decoder step
          # when prev_key and prev_value was not present in incremental_state
          if incremental_state is not None:
              self._set_input_buffer(
                  incremental_state,
                  "prev_key",
                  k.view(bsz, self.num_heads, -1, self.head_dim),
              )
              self._set_input_buffer(
                  incremental_state,
                  "prev_value",
                  v.view(bsz, self.num_heads, -1, self.head_dim),
              )
              key_padding_mask = self._get_input_buffer(
                  incremental_state, "prev_key_padding_mask"
              )
              if key_padding_mask is not None:
                  self._set_input_buffer(
                      incremental_state,
                      "prev_key_padding_mask",
                      key_padding_mask,
                  )
  
      # q.size() : bsz_X_num_heads, target_len, self.head_dim
      assert list(k.size()) == [
          bsz_X_num_heads,
          src_len,
          self.head_dim,
      ], f"key.size() :{ k.size()} [ bsz_X_num_heads, src_len, self.head_dim] : [{bsz_X_num_heads}, {src_len}, {self.head_dim}]"
  
      attn_weights = torch.bmm(q, k.transpose(1, 2))
      # attn_weights.size() : bsz_X_num_heads, target_len, src_len
  
      # Don't attend to padding symbols
      if key_padding_mask is not None:
          attn_weights = attn_weights.view(
              bsz, self.num_heads, target_len, src_len
          )
          attn_weights = attn_weights.masked_fill(
              key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
          )
          attn_weights = attn_weights.view(
              bsz_X_num_heads, target_len, src_len
          )
  
      assert list(attn_weights.size()) == [
          bsz_X_num_heads,
          target_len,
          src_len,
      ]
  
      attn_weights = F.softmax(attn_weights, dim=-1)
  
      attn_probs = self.dropout(attn_weights)
  
      attn = torch.bmm(attn_probs, v)
  
      assert list(attn.size()) == [
          bsz_X_num_heads,
          target_len,
          self.head_dim,
      ]
  
      attn = (
          attn.transpose(0, 1)
          .contiguous()
          .view(target_len, bsz, embed_dim)
      )
      attn = self.out_proj(attn)
  
      if need_weights:
          attn_weights = attn_weights.view(
              bsz, self.num_heads, target_len, src_len
          ).transpose(1, 0)
          attn_weights_out = attn_weights.mean(dim=0)
      else:
          attn_weights_out = None
  
      return attn, attn_weights_out
  
imports: from typing import Dict, Optional, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from pytext.config import ConfigBase
  from pytext.utils.usage import log_class_usage
  from torch import nn, Tensor
  from .base import PyTextIncrementalDecoderComponent
  from .utils import Linear
label_desc: Use image node_id 3 for calling the MultiheadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, key_padding_mask, need_weights, incremental_state) and returns: attn,  attn_weights_out