class_docstr: Additive Attention with linear complexity to input sequence length.
  
  Additive attention was proposed and used in FastFormer.
  See Ref. [1] for details.
  This implementation is motivated by: https://github.com/jrzaurin/pytorch-widedeep.git
  
  References:
  ----------
  [1] Wu, Chuhan, et al. "Fastformer: Additive attention can be all you need." arXiv preprint arXiv:2108.09084 (2021).
class_docstr_tok: ["Additive","Attention","with","linear","complexity","to","input","sequence","length",".","Additive","attention","was","proposed","and","used","in","FastFormer",".","See","Ref",".","[","1","]","for","details",".","This","implementation","is","motivated","by",":","https",":","\/\/github.com\/jrzaurin\/pytorch-widedeep.git","References",":","--","--","--","--","--","[","1","]","Wu",",","Chuhan",",","et","al",".","``","Fastformer",":","Additive","attention","can","be","all","you","need",".","''","arXiv","preprint","arXiv:2108.09084","(","2021",")","."]
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.
  share_qv_weights:
      if 'True', then value and query transformation parameters are shared.
  initialization:
      initialization for input projection layers. Must be one of
      :code:`['kaiming', 'xavier']`. `kaiming` is a reasonable default choice.
code: def __init__(
      self,
      *,
      d_token: int,
      n_heads: int,
      dropout: float,
      bias: bool,
      share_qv_weights: 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.
      share_qv_weights:
          if 'True', then value and query transformation parameters are shared.
      initialization:
          initialization for input projection layers. Must be one of
          :code:`['kaiming', 'xavier']`. `kaiming` is a reasonable default choice.
      """
      super().__init__()
  
      assert (
          d_token % n_heads == 0
      ), "d_token must be a multiple of n_heads"
      assert initialization in ["kaiming", "xavier"]
  
      self.head_dim = d_token // n_heads
      self.n_heads = n_heads
      self.share_qv_weights = share_qv_weights
      self.dropout = nn.Dropout(dropout)
      trainable = []
      if share_qv_weights:
          self.qv_proj = nn.Linear(d_token, d_token, bias=bias)
          trainable.extend([self.qv_proj])
      else:
          self.q_proj = nn.Linear(d_token, d_token, bias=bias)
          self.v_proj = nn.Linear(d_token, d_token, bias=bias)
          trainable.extend([self.q_proj, self.v_proj])
  
      self.k_proj = nn.Linear(d_token, d_token, bias=bias)
      self.W_q = nn.Linear(d_token, n_heads)
      self.W_k = nn.Linear(d_token, n_heads)
      self.r_out = nn.Linear(d_token, d_token)
      trainable.extend([self.k_proj, self.W_q, self.W_k, self.r_out])
  
      if initialization == "xavier":
          self.apply(init_weights)
      else:
          for m in trainable:
              if m.bias is not None:
                  nn.init.zeros_(m.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 AdditiveAttention object from inherited base classes: nn with example: obj = AdditiveAttention()