class_docstr: Multi-Head Attention module from "Attention is All You Need"
  :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.
  
  Similar to standard `dot` attention but uses
  multiple attention distributions simulataneously
  to select relevant items.
  
  .. mermaid::
  
     graph BT
        A[key]
        B[value]
        C[query]
        O[output]
        subgraph Attn
          D[Attn 1]
          E[Attn 2]
          F[Attn N]
        end
        A --> D
        C --> D
        A --> E
        C --> E
        A --> F
        C --> F
        D --> O
        E --> O
        F --> O
        B --> O
  
  Also includes several additional tricks.
  
  Args:
     head_count (int): number of parallel heads
     model_dim (int): the dimension of keys/values/queries,
         must be divisible by head_count
     dropout (float): dropout parameter
     max_relative_positions (int): max relative positions
     attn_type: "self" or "context"
code: def __init__(
      self,
      head_count: int,
      model_dim: int,
      dropout: float = 0.1,
      is_decoder: bool = True,
      max_relative_positions: int = 0,
      relative_positions_buckets: int = 0,
      rotary_interleave: bool = True,
      rotary_theta: int = 1e4,
      rotary_dim: int = 0,
      attn_type: str = None,
      self_attn_type: str = None,
      add_qkvbias=False,
      num_kv=0,
      use_ckpting=[],
      parallel_gpu=1,
  ) -> None:
      assert (
          model_dim % head_count == 0
      ), "Model dimension must be divisible by the number of heads"
      self.dim_per_head = model_dim // head_count
      super(MultiHeadedAttention, self).__init__()
      self.head_count = head_count
      self.num_kv = num_kv
      self.parallel_gpu = parallel_gpu
      if num_kv == 0:
          assert (
              model_dim % parallel_gpu == 0
          ), "Model dimension must be divisible by the number of partitions"
          self.linear_keys = skip_init(
              nn.Linear,
              in_features=model_dim,
              out_features=model_dim // parallel_gpu,
              bias=add_qkvbias,
          )
          self.linear_values = skip_init(
              nn.Linear,
              in_features=model_dim,
              out_features=model_dim // parallel_gpu,
              bias=add_qkvbias,
          )
      else:
          assert (
              self.dim_per_head * self.num_kv
          ) % parallel_gpu == 0, "Model dimension must be divisible by the number of partitions"
          self.linear_keys = skip_init(
              nn.Linear,
              in_features=model_dim,
              out_features=self.dim_per_head
              * self.num_kv
              // parallel_gpu,
              bias=add_qkvbias,
          )
          self.linear_values = skip_init(
              nn.Linear,
              in_features=model_dim,
              out_features=self.dim_per_head
              * self.num_kv
              // parallel_gpu,
              bias=add_qkvbias,
          )
      self.linear_query = skip_init(
          nn.Linear,
          in_features=model_dim,
          out_features=model_dim // parallel_gpu,
          bias=add_qkvbias,
      )
      self.softmax = nn.Softmax(dim=-1)
      self.dropout = nn.Dropout(dropout)
      self.dropout_p = dropout
      self.final_linear = skip_init(
          nn.Linear,
          in_features=model_dim // parallel_gpu,
          out_features=model_dim,
          bias=add_qkvbias,
      )
      self.is_decoder = is_decoder
      self.max_relative_positions = max_relative_positions
      self.relative_positions_buckets = relative_positions_buckets
      self.attn_type = attn_type
      self.self_attn_type = self_attn_type
      self.layer_cache = (
          False,
          {"keys": torch.tensor([]), "values": torch.tensor([])},
      )
      if relative_positions_buckets > 0:
          self.relative_attention_bias = nn.Embedding(
              relative_positions_buckets, head_count
          )
          self.relative_positions_embeddings = None
      elif max_relative_positions > 0:
          # https://arxiv.org/pdf/1803.02155.pdf
          # in the paper they suggest either two embeds
          # relative_key / relative_value or only
          # relative_key. We implemented the same embed
          # for both.
          vocab_size = max_relative_positions * 2 + 1
          self.relative_positions_embeddings = nn.Embedding(
              vocab_size, self.dim_per_head
          )
          self.relative_attention_bias = None
      else:
          self.relative_positions_embeddings = None
          self.relative_attention_bias = None
  
          if max_relative_positions == -1:  # rotary embeddings
              if rotary_dim == 0:
                  self.rotary_dim = self.dim_per_head
              else:
                  self.rotary_dim = rotary_dim
              self.rope = rotaryembeddings(
                  self.rotary_dim, base=rotary_theta
              )
              self.cos = (
                  self.rope[:, : self.rope.size(1) // 2]
                  .real.contiguous()
                  .half()
              )
              self.sin = (
                  self.rope[:, : self.rope.size(1) // 2]
                  .imag.contiguous()
                  .half()
              )
              self.rotary_interleave = rotary_interleave
              self.rotary_theta = rotary_theta
          else:
              self.cos = None
              self.sin = None
              self.rotary_interleave = None
          if max_relative_positions == -2:  # alibi positional bias
              self.alibi = AlibiPositionalBias(head_count)
  
      self.maybe_ckpt = (
          checkpoint if "mha" in use_ckpting else lambda f, x: f(x)
      )
  
      try:
          flash_pack = import_module("flash_attn")
          if (
              hasattr(flash_pack, "flash_attn_func")
              and torch.cuda.get_device_capability()[0] >= 8
          ):
              self.flash_attn_func = getattr(
                  flash_pack, "flash_attn_func"
              )
              self.flash_attn_with_kvcache = getattr(
                  flash_pack, "flash_attn_with_kvcache"
              )
              self.flash2 = True
          else:
              self.flash2 = False
      except ImportError:
          self.flash2 = False
  
imports: import torch
  import torch.nn
  from math import log, sqrt
  from torch import Tensor
  from typing import Optional, Tuple
  from torch.nn.functional import scaled_dot_product_attention
  from torch.utils.checkpoint import checkpoint
  from torch.nn.utils import skip_init
  from .alibi_position_bias import AlibiPositionalBias
  from torch.distributed import all_reduce
  from importlib import import_module
label_desc: Use image node_id 1 to create a new MultiHeadedAttention object from inherited base classes: torch.nn with example: obj = MultiHeadedAttention(head_count, model_dim, dropout, is_decoder, max_relative_positions, relative_positions_buckets, rotary_interleave, rotary_theta, rotary_dim, attn_type, _attn_type, add_qkvbias, num_kv, use_ckpting, parallel_gpu)