docstr: Pooling token mixing mechanism, as proposed in
  `Metaformer is actually what you need for vision`_, Yu et al (2021).
  
  The original notation is kept as is.
  
  .. _`Metaformer is actually what you need for vision` : https://arxiv.org/pdf/2111.11418v1.pdf
docstr_tok: ["Pooling","token","mixing","mechanism",",","as","proposed","in","`","Metaformer","is","actually","what","you","need","for","vision","`","_",",","Yu","et","al","(","2021",")",".","The","original","notation","is","kept","as","is",".","..","_","`","Metaformer","is","actually","what","you","need","for","vision","`",":","https",":","\/\/arxiv.org\/pdf\/2111.11418v1.pdf"]
code: def __init__(
      self,
      pool_size: int = 3,
      stride: int = 1,
      padding: Optional[int] = None,
      *_,
      **__,
  ):
      """
      Pooling token mixing mechanism, as proposed in
      `Metaformer is actually what you need for vision`_, Yu et al (2021).
  
      The original notation is kept as is.
  
      .. _`Metaformer is actually what you need for vision` : https://arxiv.org/pdf/2111.11418v1.pdf
      """
      super().__init__()
  
      padding = padding if padding is not None else pool_size // 2
      self.pool = nn.AvgPool2d(
          pool_size,
          stride=stride,
          padding=pool_size // 2,
          count_include_pad=False,
      )
  
      # MHA related flags:
      # kq need to have the same dimension
      self.requires_same_k_q_dimensions = False
  
      # This attention does not support attention masks
      self.supports_attention_mask = False
  
      # This "attention" (token mixing) skips the multihead attention altogether
      self.requires_skip_multi_head = True
      self.requires_input_projection = False
  
      # This operator does not really handle q,k,v
      self.requires_same_k_q_dimensions = True
  
      # This attention requires the 2d structure out of the context,
      # implictly assumed to be a squared length
      self.requires_squared_context = True
  
file_path: xformers/xformers/components/attention/pooling.py
imports: import math
  from dataclasses import dataclass
  from typing import Optional
  import torch
  import torch.nn
  from xformers.components.attention import Attention, AttentionConfig, register_attention
label_desc: Use image node_id 1 to create a new Pooling object from inherited base classes: Attention with example: obj = Pooling(pool_size, stride, padding)