class_docstr: Multihead Attention (self-/cross-) with optional 'linear' attention.
  
  To learn more about Multihead Attention, see [1]. See the implementation
  of `Transformer` and the examples below to learn how to use the compression technique
  from [2] to speed up the module when the number of tokens is large.
  
  References:
  ----------
  [1] Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding" 2018
  [2] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, Hao Ma "Linformer: Self-Attention with Linear Complexity", 2020
class_docstr_tok: ["Multihead","Attention","(","self-\/cross-",")","with","optional","'linear","'","attention",".","To","learn","more","about","Multihead","Attention",",","see","[","1","]",".","See","the","implementation","of","`","Transformer","`","and","the","examples","below","to","learn","how","to","use","the","compression","technique","from","[","2","]","to","speed","up","the","module","when","the","number","of","tokens","is","large",".","References",":","--","--","--","--","--","[","1","]","Jacob","Devlin",",","Ming-Wei","Chang",",","Kenton","Lee",",","Kristina","Toutanova","``","BERT",":","Pre-training","of","Deep","Bidirectional","Transformers","for","Language","Understanding","''","2018","[","2","]","Sinong","Wang",",","Belinda","Z.","Li",",","Madian","Khabsa",",","Han","Fang",",","Hao","Ma","``","Linformer",":","Self-Attention","with","Linear","Complexity","''",",","2020"]
docstr: Parameters
  ----------
  d_token:
      the token size. Must be a multiple of :code:`n_heads`.
  n_heads:
      the number of heads. If greater than 1, then the module will have
      an addition output layer (so called "mixing" layer).
  dropout:
      dropout rate for the attention map. The dropout is applied to
      *probabilities* and do not affect logits.
  bias:
      if `True`, then input (and output, if presented) layers also have bias.
      `True` is a reasonable default choice.
  initialization:
      initialization for input projection layers. Must be one of
      :code:`['kaiming', 'xavier']`. `kaiming` is a reasonable default choice.
  
  Raises
  ----------
      AssertionError: if requirements for the inputs are not met.
code: def __init__(
      self,
      *,
      d_token: int,
      n_heads: int,
      dropout: float,
      bias: bool,
      initialization: str,
  ) -> None:
      """
      Parameters
      ----------
      d_token:
          the token size. Must be a multiple of :code:`n_heads`.
      n_heads:
          the number of heads. If greater than 1, then the module will have
          an addition output layer (so called "mixing" layer).
      dropout:
          dropout rate for the attention map. The dropout is applied to
          *probabilities* and do not affect logits.
      bias:
          if `True`, then input (and output, if presented) layers also have bias.
          `True` is a reasonable default choice.
      initialization:
          initialization for input projection layers. Must be one of
          :code:`['kaiming', 'xavier']`. `kaiming` is a reasonable default choice.
  
      Raises
      ----------
          AssertionError: if requirements for the inputs are not met.
      """
      super().__init__()
      if n_heads > 1:
          assert (
              d_token % n_heads == 0
          ), "d_token must be a multiple of n_heads"
      assert initialization in ["kaiming", "xavier"]
  
      self.W_q = nn.Linear(d_token, d_token, bias)
      self.W_k = nn.Linear(d_token, d_token, bias)
      self.W_v = nn.Linear(d_token, d_token, bias)
      self.W_out = (
          nn.Linear(d_token, d_token, bias) if n_heads > 1 else None
      )
      self.n_heads = n_heads
      self.dropout = nn.Dropout(dropout) if dropout else None
  
      for m in [self.W_q, self.W_k, self.W_v]:
          # the "xavier" branch tries to follow torch.nn.MultiheadAttention;
          # the second condition checks if W_v plays the role of W_out; the latter one
          # is initialized with Kaiming in torch
          if initialization == "xavier" and (
              m is not self.W_v or self.W_out is not None
          ):
              # gain is needed since W_qkv is represented with 3 separate layers (it
              # implies different fan_out)
              nn.init.xavier_uniform_(m.weight, gain=1 / math.sqrt(2))
          if m.bias is not None:
              nn.init.zeros_(m.bias)
      if self.W_out is not None:
          nn.init.zeros_(self.W_out.bias)
  
imports: import enum
  import math
  import warnings
  from typing import Callable, Dict, List, Optional, Tuple, Union, cast
  import torch
  import torch.nn.functional
  from torch import Tensor, nn
  from .utils import init_weights
label_desc: Use image node_id 1 to create a new MultiheadAttention object from inherited base classes: nn with example: obj = MultiheadAttention()