class_docstr: Two-dimensional attentional pooling
  
  Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
  
  Arguments
  ---------
  spatial_dim: int
      the size of the spatial dimension
  embed_dim: int
      the embedding dimension
  num_heads_channels: int
      the number of attention heads
  output_dim: int
      the output dimension
  
  Example
  -------
  >>> attn_pool = AttentionPool2d(
  ...     spatial_dim=64,
  ...     embed_dim=16,
  ...     num_heads_channels=2,
  ...     output_dim=4
  ... )
  >>> x = torch.randn(4, 1, 64, 64)
  >>> x_pool = attn_pool(x)
  >>> x_pool.shape
  torch.Size([4, 4])
class_docstr_tok: ["Two-dimensional","attentional","pooling","Adapted","from","CLIP",":","https",":","\/\/github.com\/openai\/CLIP\/blob\/main\/clip\/model.py","Arguments","--","--","--","--","-","spatial_dim",":","int","the","size","of","the","spatial","dimension","embed_dim",":","int","the","embedding","dimension","num_heads_channels",":","int","the","number","of","attention","heads","output_dim",":","int","the","output","dimension","Example","--","--","--","-",">",">",">","attn_pool","=","AttentionPool2d","(","...","spatial_dim=64",",","...","embed_dim=16",",","...","num_heads_channels=2",",","...","output_dim=4","...",")",">",">",">","x","=","torch.randn","(","4",",","1",",","64",",","64",")",">",">",">","x_pool","=","attn_pool","(","x",")",">",">",">","x_pool.shape","torch.Size","(","[","4",",","4","]",")"]
code: def __init__(
      self,
      spatial_dim: int,
      embed_dim: int,
      num_heads_channels: int,
      output_dim: int = None,
  ):
      super().__init__()
      self.positional_embedding = nn.Parameter(
          torch.randn(embed_dim, spatial_dim**2 + 1)
          / embed_dim**0.5
      )
      self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
      self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
      self.num_heads = embed_dim // num_heads_channels
      self.attention = QKVAttention(self.num_heads)
  
imports: from abc import abstractmethod
  from speechbrain.utils.data_utils import pad_divisible
  from .autoencoders import NormalizingAutoencoder
  import math
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new AttentionPool2d object from inherited base classes: nn with example: obj = AttentionPool2d(spatial_dim, embed_dim, num_heads_channels, output_dim)