class_docstr: A multi-head masked self-attention dispatch mechanism, with a projection at the end,
  following the architecture proposed in `Attention is all you need`_, Vaswani et al.
  
  The actual attention mechanism can vary, as well as the projections.
  This can be used to wrap the proposed attention mechanisms and make them multi-head aware,
  but it is optional.
  
  Args:
      dim_model: The model/embedding dimension
      num_heads: The number of heads being used
      attention: The attention mechanism (needs to be registered to the xformers library)
      bias: Whether to use bias for the projections : (Q, K, V, Output)
      residual_dropout: Amount of dropout on the residual path
      use_separate_proj_weight: Use different weights for the Q, K, V projections
      dim_key: Optionally use a different dimension for the key
      dim_value:  Optionally use a different dimension for the value
      in_proj_container: Optionally provide the input projection module
      use_rotary_embeddings: Use rotary embeddings
      out_proj: Optionally provide the output projection module
  
  
  .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762v5
class_docstr_tok: ["A","multi-head","masked","self-attention","dispatch","mechanism",",","with","a","projection","at","the","end",",","following","the","architecture","proposed","in","`","Attention","is","all","you","need","`","_",",","Vaswani","et","al",".","The","actual","attention","mechanism","can","vary",",","as","well","as","the","projections",".","This","can","be","used","to","wrap","the","proposed","attention","mechanisms","and","make","them","multi-head","aware",",","but","it","is","optional",".","Args",":","dim_model",":","The","model\/embedding","dimension","num_heads",":","The","number","of","heads","being","used","attention",":","The","attention","mechanism","(","needs","to","be","registered","to","the","xformers","library",")","bias",":","Whether","to","use","bias","for","the","projections",":","(","Q",",","K",",","V",",","Output",")","residual_dropout",":","Amount","of","dropout","on","the","residual","path","use_separate_proj_weight",":","Use","different","weights","for","the","Q",",","K",",","V","projections","dim_key",":","Optionally","use","a","different","dimension","for","the","key","dim_value",":","Optionally","use","a","different","dimension","for","the","value","in_proj_container",":","Optionally","provide","the","input","projection","module","use_rotary_embeddings",":","Use","rotary","embeddings","out_proj",":","Optionally","provide","the","output","projection","module","..","_","`","Attention","is","all","you","need","`",":","https",":","\/\/arxiv.org\/abs\/1706.03762v5"]
code: def __init__(
      self,
      dim_model: int,
      num_heads: int,
      attention: Attention,
      bias: Tuple[bool, bool, bool, bool] = (True, True, True, True),
      residual_dropout: float = 0.0,
      use_separate_proj_weight: bool = True,
      dim_key: Optional[int] = None,
      dim_value: Optional[int] = None,
      in_proj_container: Optional[InputProjection] = None,
      use_rotary_embeddings: Optional[bool] = False,
      out_proj: Optional[nn.Module] = None,
      *args,
      **kwargs,
  ):
      super().__init__()
  
      if isinstance(bias, bool):
          logger.warning(
              "Single bias value provided for the MHA projections."
              + f" Assuming the same parameter ({bias}) is to be used everywhere"
          )
          bias = (bias, bias, bias, bias)
  
      assert (
          dim_model % num_heads == 0
      )  # static preset for now, each head works on 1/d the embeddings, could be relaxed
      assert num_heads > 0
  
      # Popular default is that all latent dimensions are the same
      dim_key, dim_value = map(
          lambda x: x if x else dim_model, (dim_key, dim_value)
      )
  
      self.num_heads = num_heads
      self.dim_key_head = dim_key // num_heads
      self.dim_value_head = dim_value // num_heads
      self.dim_model = dim_model
      self.attention = attention
  
      # key, query, value projections for all heads
      # critical options are
      # - are we sharing weights ?
      # - are we adding biases ?
      if attention.requires_input_projection:
          self.in_proj_container = (
              in_proj_container
              if in_proj_container is not None
              else InputProjection(
                  query_proj_params=InputProjectionConfig(
                      dim_model, dim_key, bias=bias[0]
                  ),
                  key_proj_params=InputProjectionConfig(
                      dim_model, dim_key, bias=bias[1]
                  ),
                  value_proj_params=InputProjectionConfig(
                      dim_model, dim_value, bias=bias[2]
                  ),
                  use_separate_proj_weight=use_separate_proj_weight,
              )
          )
  
      # Optional rotary embeddings
      self.rotary_embeddings = (
          RotaryEmbedding(self.dim_key_head)
          if use_rotary_embeddings
          else None
      )
  
      # Regularization
      self.resid_drop = nn.Dropout(residual_dropout, inplace=False)
  
      # Output projection
      self.proj = (
          out_proj
          if out_proj
          else nn.Linear(dim_model, dim_model, bias=bias[3])
      )
      if (
          isinstance(self.proj, nn.Linear)
          and self.proj.bias is not None
      ):
          constant_(self.proj.bias, 0.0)
  
file_path: xformers/xformers/components/multi_head_dispatch.py
imports: import logging
  from dataclasses import asdict, dataclass
  from typing import Optional, Tuple
  import torch
  import torch.nn
  from torch.nn.init import constant_
  from xformers.components.attention import Attention
  from xformers.components.input_projection import InputProjection, InputProjectionConfig
  from xformers.components.positional_embedding import RotaryEmbedding
label_desc: Use image node_id 1 to create a new MultiHeadDispatch object from inherited base classes: nn with example: obj = MultiHeadDispatch(dim_model, num_heads, attention, bias, residual_dropout, use_separate_proj_weight, dim_key, dim_value, in_proj_container, use_rotary_embeddings, out_proj)