class_docstr: This class implements attention pooling layer, as
  described in `Clip <https://arxiv.org/pdf/2103.00020.pdf>`_, and should be
  used for CNN-style models, including MobileViTs
class_docstr_tok: ["This","class","implements","attention","pooling","layer",",","as","described","in","`","Clip","<","https",":","\/\/arxiv.org\/pdf\/2103.00020.pdf",">","`","_",",","and","should","be","used","for","CNN-style","models",",","including","MobileViTs"]
code: def forward(self, x: Tensor, *args, **kwargs) -> Tensor:
      assert (
          x.dim() == 4
      ), "Input should be 4-dimensional (Batch, in_channels, height, width). Got: {}".format(
          x.shape
      )
      # x is [batch, in_channels, height, width]
      # For CNN-style architectures, including MobileViTs
      batch_size, in_channels, in_height, in_width = x.shape
  
      # Flatten the feature map
      # [batch, in_channels, height, width] --> [batch, in_channels, height*width]
      x = x.reshape(batch_size, in_channels, in_height * in_width)
  
      if self.use_pytorch_mha:
          # we need sequence first.
          # [batch, in_channels, height*width] --> [height*width, batch, in_channels]
          x = x.permute(2, 0, 1)
  
          # global pool
          # [height*width, batch, in_channels] --> [1, batch, in_channels]
          global_token = torch.mean(x, dim=0, keepdim=True)
  
          num_pixels = x.shape[0]
  
          # add positional embedding to pixels
          pos_emb = self.positional_embedding(num_pixels).to(
              device=x.device, dtype=x.dtype
          )
          x = x + pos_emb
  
          # concat the global token with pixel tokens
          # [1, batch, in_channels] || [height*width, batch, in_channels] --> [1 + height*width, batch, in_channels]
          x = torch.cat([global_token, x], dim=0)
  
          # do attention
          x = self.multi_head_attn(x, use_pytorch_mha=True)
  
          # extract embeddings corresponding to global token
          x = x[0]
      else:
          # [batch, in_channels, height*width] --> # [batch, height*width, in_channels]
          x = x.transpose(1, 2)
  
          # global pool
          # [batch, height*width, in_channels] --> [batch, 1, in_channels]
          global_token = torch.mean(x, dim=1, keepdim=True)
  
          num_pixels = x.shape[1]
          # add positional embedding to pixels
          pos_emb = self.positional_embedding(num_pixels).to(
              device=x.device, dtype=x.dtype
          )
          x = x + pos_emb
  
          # concat the global token with pixel tokens
          # [batch, 1, in_channels] || [batch, height*width, in_channels] --> [batch, 1 + height*width, in_channels]
          x = torch.cat([global_token, x], dim=1)
  
          # do attention
          x = self.multi_head_attn(x, use_pytorch_mha=False)
  
          # extract embeddings corresponding to global token
          x = x[:, 0]
  
      if self.normalize_features:
          x = F.normalize(x, dim=-1)
      return x
  
file_path: ml-cvnets/cvnets/image_projection_layers/attention_pool_2d.py
imports: import argparse
  import torch
  from torch import Tensor, nn
  from torch.nn import functional
  from cvnets.image_projection_layers import IMAGE_PROJECTION_HEAD_REGISTRY, BaseImageProjectionHead
  from cvnets.layers import MultiHeadAttention, PositionalEmbedding
  from utils import logger
label_desc: Use image node_id 4 for calling the AttentionPool2dHead obj's underlying member method code with example usage: obj.forward(x) and returns: x