class_docstr: A MLP block that contains two linear layers with a normalization layer.
  the MLP block is used in a transformer model after the attention block.
  ::
                                      Input
                                         ↓
                                      Reshape
                                         ↓
                                        Pool
                                         ↓
                                      Reshape
                                         ↓
                                        norm
  Args:
      thw_shape (`List`):
          the shape of the input tensor (before flattening).
      pool (`Callable`):
          Pool operation that is applied to the input tensor.
          If pool is None, return the input tensor.
      has_cls_embed (`bool`):
          whether the input tensor contains cls token. Pool operation excludes cls token.
      norm (`Callable`):
          Optional normalization operation applied to tensor after pool.
  Returns:
      tensor (`torch.Tensor`):
          Input tensor after pool.
      thw_shape (`List[int]`):
          Output tensor shape (before flattening).
code: def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[int]]:
      if self.pool is None:
          return x, self.thw_shape
      tensor_dim = x.ndim
      if tensor_dim == 4:
          pass
      elif tensor_dim == 3:
          x = x.unsqueeze(1)
      else:
          raise NotImplementedError(
              f"Unsupported input dimension {x.shape}"
          )
  
      if self.has_cls_embed:
          cls_tok, x = x[:, :, :1, :], x[:, :, 1:, :]
  
      b, n, _, c = x.shape
      t, h, w = self.thw_shape
      x = (
          x.reshape(b * n, t, h, w, c)
          .permute(0, 4, 1, 2, 3)
          .contiguous()
      )
  
      x = self.pool(x)
  
      thw_shape = [x.shape[2], x.shape[3], x.shape[4]]
      l_pooled = x.shape[2] * x.shape[3] * x.shape[4]
      x = x.reshape(b, n, c, l_pooled).transpose(2, 3)
      if self.has_cls_embed:
          x = torch.cat((cls_tok, x), dim=2)
      if self.norm is not None:
          x = self.norm(x)
  
      if tensor_dim == 4:
          pass
      else:  # For the case tensor_dim == 3.
          x = x.squeeze(1)
      return x, thw_shape
  
imports: import torch
  from torch import nn
  from typing import List, Tuple
label_desc: Use image node_id 2 for calling the AttentionPool obj's underlying member method code with example usage: obj.forward(x) and returns: x,  thw_shape, x,  self