docstr: Kernelized attention, as proposed in Performers_
  ("Rethinking attention with performers." K. Choromanski et al. (2020).).
  
  FAVOR stands for "Fast Attention Via positive Orthogonal Random features"
  
  Args:
      dropout (float): the probability of an output to be randomly dropped at training time
      dim_features (int): the dimension of the random features space
      iter_before_redraw (int): the number of steps (forward calls) before a redraw of the features
      feature_map_type (FeatureMapType): the type of feature map being used,
      for instance orthogonal random features.
  
  .. _Performers: https://arxiv.org/pdf/2009.14794v1.pdf
docstr_tok: ["Kernelized","attention",",","as","proposed","in","Performers_","(","``","Rethinking","attention","with","performers",".","''","K.","Choromanski","et","al",".","(","2020",")",".",")",".","FAVOR","stands","for","``","Fast","Attention","Via","positive","Orthogonal","Random","features","''","Args",":","dropout","(","float",")",":","the","probability","of","an","output","to","be","randomly","dropped","at","training","time","dim_features","(","int",")",":","the","dimension","of","the","random","features","space","iter_before_redraw","(","int",")",":","the","number","of","steps","(","forward","calls",")","before","a","redraw","of","the","features","feature_map_type","(","FeatureMapType",")",":","the","type","of","feature","map","being","used",",","for","instance","orthogonal","random","features",".","..","_Performers",":","https",":","\/\/arxiv.org\/pdf\/2009.14794v1.pdf"]
code: def __init__(
      self,
      causal: bool = False,
      dropout: float = 0.0,
      dim_features: Optional[int] = None,
      dim_head: Optional[int] = None,
      iter_before_redraw: Optional[int] = None,
      feature_map_type: FeatureMapType = FeatureMapType.SMReg,
      normalize_inputs: bool = False,
      *_,
      **__,
  ):
      r"""
      Kernelized attention, as proposed in Performers_
      ("Rethinking attention with performers." K. Choromanski et al. (2020).).
  
      FAVOR stands for "Fast Attention Via positive Orthogonal Random features"
  
      Args:
          dropout (float): the probability of an output to be randomly dropped at training time
          dim_features (int): the dimension of the random features space
          iter_before_redraw (int): the number of steps (forward calls) before a redraw of the features
          feature_map_type (FeatureMapType): the type of feature map being used,
          for instance orthogonal random features.
  
      .. _Performers: https://arxiv.org/pdf/2009.14794v1.pdf
      """
      super().__init__()
  
      self.causal = causal
      self.iter_before_redraw = (
          (2 * iter_before_redraw)
          if iter_before_redraw is not None
          else iter_before_redraw
      )  # This will be used for both key and query
      self.normalize_inputs = normalize_inputs
      self.feature_map_type = feature_map_type
      self.attn_drop = nn.Dropout(dropout, inplace=True)
  
      # Setup dimension-dependent variables
      # Reasonable dimension default
      if dim_features is None:
          assert (
              dim_head is not None
          ), "dim_features or dim_head needs to be passed"
          self.dim_features = math.ceil(
              dim_head * (1 + math.log2(dim_head))
          )
          self.dim_features = 2 * (
              self.dim_features // 2
          )  # needs to be even for some variants
          logger.info(
              f"FAVOR: Automatically setting the random mapping dimension to {self.dim_features} from {dim_head}"
          )
      else:
          self.dim_features = dim_features
  
      feature_map_constructor = {
          FeatureMapType.SMHyp: SMHyperbolic,
          FeatureMapType.SMReg: SMReg,
          FeatureMapType.SMOrf: SMOrf,
      }[self.feature_map_type]
  
      feature_settings = {
          "dim_features": self.dim_features,
          "iter_before_redraw": self.iter_before_redraw,
          "normalize_inputs": self.normalize_inputs,
      }
  
      self.feature_map: FeatureMap = feature_map_constructor(**feature_settings)  # type: ignore
  
      # Properties specific to this attention mechanism
      self.supports_attention_mask = False
      self.supports_key_padding_mask = False
  
file_path: xformers/xformers/components/attention/favor.py
imports: import logging
  import math
  from dataclasses import dataclass
  from typing import Optional, Tuple
  import torch
  import torch.nn
  from torch.cuda.amp import autocast
  from xformers.components.attention import Attention, AttentionConfig, register_attention
  from xformers.components.attention.feature_maps import FeatureMap, FeatureMapType, SMHyperbolic, SMOrf, SMReg
label_desc: Use image node_id 1 to create a new FavorAttention object from inherited base classes: Attention with example: obj = FavorAttention(causal, dropout, dim_features, dim_head, iter_before_redraw, feature_map_type, normalize_inputs)