class_docstr: Multiheaded Scaled Dot Product Attention. This function
  has the same exact signature as the one used in pytorch_translate
  with the added benefit of supporting torchscript
class_docstr_tok: ["Multiheaded","Scaled","Dot","Product","Attention",".","This","function","has","the","same","exact","signature","as","the","one","used","in","pytorch_translate","with","the","added","benefit","of","supporting","torchscript"]
docstr: Computes MultiheadAttention with respect to either a vector
  or a tensor
  
  Inputs:
      decoder_state: (bsz x decoder_hidden_state_dim) or
          (bsz x T x decoder_hidden_state_dim)
      source_hids: srclen x bsz x context_dim
      src_lengths: bsz x 1, actual sequence lengths
      squeeze: Whether or not to squeeze on the time dimension.
          Even if decoder_state.dim() is 2 dimensional an
          explicit time step dimension will be unsqueezed.
  Outputs:
    [batch_size, max_src_len] if decoder_state.dim() == 2 & squeeze
      or
    [batch_size, 1, max_src_len] if decoder_state.dim() == 2 & !squeeze
      or
    [batch_size, T, max_src_len] if decoder_state.dim() == 3 & !squeeze
      or
    [batch_size, T, max_src_len] if decoder_state.dim() == 3 & squeeze & T != 1
      or
    [batch_size, max_src_len] if decoder_state.dim() == 3 & squeeze & T == 1
docstr_tok: ["Computes","MultiheadAttention","with","respect","to","either","a","vector","or","a","tensor","Inputs",":","decoder_state",":","(","bsz","x","decoder_hidden_state_dim",")","or","(","bsz","x","T","x","decoder_hidden_state_dim",")","source_hids",":","srclen","x","bsz","x","context_dim","src_lengths",":","bsz","x","1",",","actual","sequence","lengths","squeeze",":","Whether","or","not","to","squeeze","on","the","time","dimension",".","Even","if","decoder_state.dim","(",")","is","2","dimensional","an","explicit","time","step","dimension","will","be","unsqueezed",".","Outputs",":","[","batch_size",",","max_src_len","]","if","decoder_state.dim","(",")","==","2","&","squeeze","or","[","batch_size",",","1",",","max_src_len","]","if","decoder_state.dim","(",")","==","2","&","!","squeeze","or","[","batch_size",",","T",",","max_src_len","]","if","decoder_state.dim","(",")","==","3","&","!","squeeze","or","[","batch_size",",","T",",","max_src_len","]","if","decoder_state.dim","(",")","==","3","&","squeeze","&","T","!","=","1","or","[","batch_size",",","max_src_len","]","if","decoder_state.dim","(",")","==","3","&","squeeze","&","T","==","1"]
code: def forward(
      self,
      decoder_state: Tensor,
      source_hids: Tensor,
      src_len_mask: Optional[Tensor],
      squeeze: bool = True,
  ) -> Tuple[Tensor, Tensor]:
      """
      Computes MultiheadAttention with respect to either a vector
      or a tensor
  
      Inputs:
          decoder_state: (bsz x decoder_hidden_state_dim) or
              (bsz x T x decoder_hidden_state_dim)
          source_hids: srclen x bsz x context_dim
          src_lengths: bsz x 1, actual sequence lengths
          squeeze: Whether or not to squeeze on the time dimension.
              Even if decoder_state.dim() is 2 dimensional an
              explicit time step dimension will be unsqueezed.
      Outputs:
        [batch_size, max_src_len] if decoder_state.dim() == 2 & squeeze
          or
        [batch_size, 1, max_src_len] if decoder_state.dim() == 2 & !squeeze
          or
        [batch_size, T, max_src_len] if decoder_state.dim() == 3 & !squeeze
          or
        [batch_size, T, max_src_len] if decoder_state.dim() == 3 & squeeze & T != 1
          or
        [batch_size, max_src_len] if decoder_state.dim() == 3 & squeeze & T == 1
      """
      if decoder_state.dim() == 3:
          query = decoder_state
      elif decoder_state.dim() == 2:
          query = decoder_state.unsqueeze(1)
      else:
          raise ValueError(
              "decoder state must be either 2 or 3 dimensional"
          )
      query = query.transpose(0, 1)
      value = key = source_hids
  
      attn, attn_weights = self._attn.forward(
          query,
          key,
          value,
          key_padding_mask=src_len_mask,
          need_weights=True,
      )
      # Need to satify torchscript here
      if attn_weights is None:
          raise NotImplementedError("")
      # attn.shape = T X bsz X embed_dim
      # attn_weights.shape = bsz X T X src_len
  
      attn_weights = attn_weights.transpose(0, 2)
      # attn_weights.shape = src_len X T X bsz
  
      if squeeze:
          attn = attn.squeeze(0)
          # attn.shape = squeeze(T) X bsz X embed_dim
          attn_weights = attn_weights.squeeze(1)
          # attn_weights.shape = src_len X squeeze(T) X bsz
          return attn, attn_weights
      return attn, attn_weights
  
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 2 for calling the DecoupledMultiheadAttention obj's underlying member method code with example usage: obj.forward(decoder_state, source_hids, src_len_mask, squeeze) and returns: attn,  attn_weights, attn,  attn_weights