class_docstr: Multi-headed attention.
  
  See "Attention Is All You Need" for more details.
class_docstr_tok: ["Multi-headed","attention",".","See","``","Attention","Is","All","You","Need","''","for","more","details","."]
docstr: Input shape: Time x Batch x Channel
  
  Self-attention can be implemented by passing in the same arguments for
  query, key and value. Future timesteps can be masked with the
  `mask_future_timesteps` argument. Padding elements can be excluded from
  the key by passing a binary ByteTensor (`key_padding_mask`) with shape:
  batch x src_len, where padding elements are indicated by 1s.
docstr_tok: ["Input","shape",":","Time","x","Batch","x","Channel","Self-attention","can","be","implemented","by","passing","in","the","same","arguments","for","query",",","key","and","value",".","Future","timesteps","can","be","masked","with","the","`","mask_future_timesteps","`","argument",".","Padding","elements","can","be","excluded","from","the","key","by","passing","a","binary","ByteTensor","(","`","key_padding_mask","`",")","with","shape",":","batch","x","src_len",",","where","padding","elements","are","indicated","by","1s","."]
code: def forward(
      self,
      query,
      key,
      value,
      key_padding_mask=None,
      need_weights=False,
      attn_mask=None,
      is_training=True,
  ):
      """Input shape: Time x Batch x Channel
  
      Self-attention can be implemented by passing in the same arguments for
      query, key and value. Future timesteps can be masked with the
      `mask_future_timesteps` argument. Padding elements can be excluded from
      the key by passing a binary ByteTensor (`key_padding_mask`) with shape:
      batch x src_len, where padding elements are indicated by 1s.
      """
      if self.separate_qkv_params:
          input_weights = (
              torch.cat(
                  [
                      self.q_weight.view(
                          self.num_heads,
                          1,
                          self.head_dim,
                          self.embed_dim,
                      ),
                      self.k_weight.view(
                          self.num_heads,
                          1,
                          self.head_dim,
                          self.embed_dim,
                      ),
                      self.v_weight.view(
                          self.num_heads,
                          1,
                          self.head_dim,
                          self.embed_dim,
                      ),
                  ],
                  dim=1,
              )
              .reshape(3 * self.embed_dim, self.embed_dim)
              .contiguous()
          )
      else:
          input_weights = self.in_proj_weight
      if self.bias:
          if self.separate_qkv_params:
              input_bias = (
                  torch.cat(
                      [
                          self.q_bias.view(
                              self.num_heads, 1, self.head_dim
                          ),
                          self.k_bias.view(
                              self.num_heads, 1, self.head_dim
                          ),
                          self.v_bias.view(
                              self.num_heads, 1, self.head_dim
                          ),
                      ],
                      dim=1,
                  )
                  .reshape(3 * self.embed_dim)
                  .contiguous()
              )
          else:
              input_bias = self.in_proj_bias
      else:
          input_bias = None
      if key_padding_mask is not None:
          assert (
              attn_mask is None
          ), "ERROR attn_mask and key_padding_mask should not be both defined!"
          mask = key_padding_mask
      elif attn_mask is not None:
          assert (
              self.mask_additive == False
          ), "additive mask not supported for time mask"
          mask = attn_mask
      else:
          mask = None
  
      if self.include_norm_add:
          if self.impl == "fast":
              outputs = self.attn_func(
                  attn_mask is not None,
                  is_training,
                  self.num_heads,
                  query,
                  self.lyr_nrm_gamma_weights,
                  self.lyr_nrm_beta_weights,
                  input_weights,
                  self.out_proj_weight,
                  mask,
                  self.dropout,
              )
          else:
              lyr_nrm_results = self.lyr_nrm(query)
              outputs = self.attn_func(
                  attn_mask is not None,
                  is_training,
                  self.num_heads,
                  self.scaling,
                  lyr_nrm_results,
                  input_weights,
                  self.out_proj_weight,
                  input_bias,
                  self.out_proj_bias,
                  mask,
                  self.mask_additive,
                  self.dropout,
              )
              if is_training:
                  outputs = jit_dropout_add(
                      outputs, query, self.dropout, is_training
                  )
              else:
                  outputs = outputs + query
      else:
          if self.impl == "fast":
              outputs = self.attn_func(
                  attn_mask is not None,
                  is_training,
                  self.num_heads,
                  query,
                  input_weights,
                  self.out_proj_weight,
                  input_bias,
                  self.out_proj_bias,
                  mask,
                  self.mask_additive,
                  self.dropout,
              )
          else:
              outputs = self.attn_func(
                  attn_mask is not None,
                  is_training,
                  self.num_heads,
                  self.scaling,
                  query,
                  input_weights,
                  self.out_proj_weight,
                  input_bias,
                  self.out_proj_bias,
                  mask,
                  self.mask_additive,
                  self.dropout,
              )
  
      return outputs, None
  
file_path: apex/apex/contrib/multihead_attn/self_multihead_attn.py
imports: import math
  import torch
  from torch import nn
  from torch.nn import Parameter
  import torch.nn.functional
  from .self_multihead_attn_func import self_attn_func
  from .fast_self_multihead_attn_func import fast_self_attn_func
  from .fast_self_multihead_attn_norm_add_func import fast_self_attn_norm_add_func
  from apex.normalization.fused_layer_norm import FusedLayerNorm
label_desc: Use image node_id 3 for calling the SelfMultiheadAttn obj's underlying member method code with example usage: obj.forward(query, key, value, key_padding_mask, need_weights, attn_mask, is_training) and returns: outputs,  None