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","."]
code: def __init__(
      self,
      embed_dim,
      num_heads,
      dropout=0.0,
      bias=False,
      include_norm_add=False,
      impl="fast",
  ):
      super().__init__()
      self.embed_dim = embed_dim
      self.num_heads = num_heads
      self.dropout = dropout
      self.head_dim = embed_dim // num_heads
      assert (
          self.head_dim * num_heads == self.embed_dim
      ), "embed_dim must be divisible by num_heads"
      self.bias = bias
      self.include_norm_add = include_norm_add
      self.impl = impl
      self.scaling = self.head_dim**-0.5
  
      self.in_proj_weight_q = Parameter(
          torch.empty(embed_dim, embed_dim)
      )
      self.in_proj_weight_kv = Parameter(
          torch.empty(2 * embed_dim, embed_dim)
      )
      self.out_proj_weight = Parameter(
          torch.empty(embed_dim, embed_dim)
      )
      if self.bias:
          assert (
              impl != "fast"
          ), "ERROR! The Fast implementation does not support biases!"
          self.in_proj_bias_q = Parameter(torch.empty(embed_dim))
          self.in_proj_bias_kv = Parameter(torch.empty(2 * embed_dim))
          self.out_proj_bias = Parameter(torch.empty(embed_dim))
      else:
          self.register_parameter("in_proj_bias_q", None)
          self.register_parameter("in_proj_bias_kv", None)
          self.in_proj_bias_q = None
          self.in_proj_bias_kv = None
          self.out_proj_bias = None
      if self.include_norm_add:
          if impl == "fast":
              self.lyr_nrm_gamma_weights = Parameter(
                  torch.empty(embed_dim)
              )
              self.lyr_nrm_beta_weights = Parameter(
                  torch.empty(embed_dim)
              )
              self.lyr_nrm = None
          else:
              self.register_parameter("lyr_norm_gamma_weights", None)
              self.register_parameter("lyr_norm_beta_weights", None)
              self.lyr_nrm_gamma_weights = None
              self.lyr_nrm_beta_weights = None
              self.lyr_nrm = FusedLayerNorm(embed_dim)
      self.reset_parameters()
  
      if self.include_norm_add:
          if impl == "fast":
              self.attn_func = fast_encdec_attn_norm_add_func
          elif impl == "default":
              self.attn_func = encdec_attn_func
          else:
              assert False, "Unsupported impl: {} !".format(impl)
      else:
          if impl == "fast":
              self.attn_func = fast_encdec_attn_func
          elif impl == "default":
              self.attn_func = encdec_attn_func
          else:
              assert False, "Unsupported impl: {} !".format(impl)
  
file_path: apex/apex/contrib/multihead_attn/encdec_multihead_attn.py
imports: import math
  import torch
  from torch import nn
  from torch.nn import Parameter
  import torch.nn.functional
  from .encdec_multihead_attn_func import encdec_attn_func
  from .fast_encdec_multihead_attn_func import fast_encdec_attn_func
  from .fast_encdec_multihead_attn_norm_add_func import fast_encdec_attn_norm_add_func
  from apex.normalization.fused_layer_norm import FusedLayerNorm
label_desc: Use image node_id 1 to create a new EncdecMultiheadAttn object from inherited base classes: nn with example: obj = EncdecMultiheadAttn(embed_dim, num_heads, dropout, bias, include_norm_add, impl)