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[int]`):
          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).
class_docstr_tok: ["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","\u2193","Reshape","\u2193","Pool","\u2193","Reshape","\u2193","norm","Args",":","thw_shape","(","`","List","[","int","]","`",")",":","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, thw_shape
  ) -> Tuple[torch.Tensor, List[int]]:
      if self.pool is None:
          return x, thw_shape
      self.tensor_dim = x.ndim
      if self.tensor_dim == 4:
          pass
      elif self.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 = thw_shape
      x = rearrange(
          x,
          "b n (t h w) c -> (b n) c t h w",
          b=b,
          n=n,
          c=c,
          t=t,
          h=h,
          w=w,
      )
  
      x = self.pool(x)
      out_thw_shape = [x.shape[2], x.shape[3], x.shape[4]]
      x = rearrange(x, "(b n) c t h w  -> b n (t h w) c", b=b, n=n, c=c)
      if self.has_cls_embed:
          x = torch.cat((cls_tok, x), dim=2)
      if self.norm is not None:
          x = self.norm(x)
  
      if self.tensor_dim == 4:
          pass
      else:
          x = x.squeeze(1)
      return x, out_thw_shape
  
imports: import math
  from functools import partial
  from typing import Tuple, List
  import numpy
  import torch
  from torch import nn
  from torch.nn.init import trunc_normal_
  from einops import rearrange
  from towhee.models.layers.layers_with_relprop import GELU, Linear, Dropout, Conv2d, Conv3d, Softmax, Sigmoid, LayerNorm, Einsum, Add, Clone, MaxPool3d, IndexSelect
label_desc: Use image node_id 2 for calling the AttentionPool obj's underlying member method code with example usage: obj.forward(x, thw_shape) and returns: x,  out_thw_shape, x,  thw_shape