class_docstr: MaxViT Transformer block.
      With block partition:
      x ← x + Unblock(RelAttention(Block(LN(x))))
      x ← x + MLP(LN(x))
      With grid partition:
      x ← x + Ungrid(RelAttention(Grid(LN(x))))
      x ← x + MLP(LN(x))
      Layer Normalization (LN) is applied after the grid/window partition to prevent multiple reshaping operations.
      Grid/window reverse (Unblock/Ungrid) is performed on the final output for the same reason.
  Args:
      in_channels (`int`):
          Number of input channels.
      partition_function (`Callable`):
          Partition function to be utilized (grid or window partition).
      reverse_function (`Callable`):
          Reverse function to be utilized  (grid or window reverse).
      num_heads (`int`):
          Number of attention heads. Default 32
      grid_window_size (`Tuple[int, int]`):
          Grid/Window size to be utilized. Default (7, 7)
      attn_drop (`float`):
          Dropout ratio of attention weight. Default: 0.0
      drop (`float`):
          Dropout ratio of output. Default: 0.0
      drop_path (`float`):
          Dropout ratio of path. Default: 0.0
      mlp_ratio (`float`):
          Ratio of mlp hidden dim to embedding dim. Default: 4.0
      act_layer (`nn.Module`):
          Type of activation layer to be utilized. Default: nn.GELU
      norm_layer (`nn.Module`):
          Type of normalization layer to be utilized. Default: nn.BatchNorm2d
class_docstr_tok: ["MaxViT","Transformer","block",".","With","block","partition",":","x","\u2190","x","+","Unblock","(","RelAttention","(","Block","(","LN","(","x",")",")",")",")","x","\u2190","x","+","MLP","(","LN","(","x",")",")","With","grid","partition",":","x","\u2190","x","+","Ungrid","(","RelAttention","(","Grid","(","LN","(","x",")",")",")",")","x","\u2190","x","+","MLP","(","LN","(","x",")",")","Layer","Normalization","(","LN",")","is","applied","after","the","grid\/window","partition","to","prevent","multiple","reshaping","operations",".","Grid\/window","reverse","(","Unblock\/Ungrid",")","is","performed","on","the","final","output","for","the","same","reason",".","Args",":","in_channels","(","`","int","`",")",":","Number","of","input","channels",".","partition_function","(","`","Callable","`",")",":","Partition","function","to","be","utilized","(","grid","or","window","partition",")",".","reverse_function","(","`","Callable","`",")",":","Reverse","function","to","be","utilized","(","grid","or","window","reverse",")",".","num_heads","(","`","int","`",")",":","Number","of","attention","heads",".","Default","32","grid_window_size","(","`","Tuple","[","int",",","int","]","`",")",":","Grid\/Window","size","to","be","utilized",".","Default","(","7",",","7",")","attn_drop","(","`","float","`",")",":","Dropout","ratio","of","attention","weight",".","Default",":","0.0","drop","(","`","float","`",")",":","Dropout","ratio","of","output",".","Default",":","0.0","drop_path","(","`","float","`",")",":","Dropout","ratio","of","path",".","Default",":","0.0","mlp_ratio","(","`","float","`",")",":","Ratio","of","mlp","hidden","dim","to","embedding","dim",".","Default",":","4.0","act_layer","(","`","nn.Module","`",")",":","Type","of","activation","layer","to","be","utilized",".","Default",":","nn.GELU","norm_layer","(","`","nn.Module","`",")",":","Type","of","normalization","layer","to","be","utilized",".","Default",":","nn.BatchNorm2d"]
docstr: Forward pass.
  Args:
      data (torch.Tensor): Input tensor of the shape [B, C_in, H, W].
  Returns:
      output (torch.Tensor): Output tensor of the shape [B, C_out, H (// 2), W (// 2)].
docstr_tok: ["Forward","pass",".","Args",":","data","(","torch.Tensor",")",":","Input","tensor","of","the","shape","[","B",",","C_in",",","H",",","W","]",".","Returns",":","output","(","torch.Tensor",")",":","Output","tensor","of","the","shape","[","B",",","C_out",",","H","(","\/\/","2",")",",","W","(","\/\/","2",")","]","."]
code: def forward(self, data: torch.Tensor) -> torch.Tensor:
      """Forward pass.
      Args:
          data (torch.Tensor): Input tensor of the shape [B, C_in, H, W].
      Returns:
          output (torch.Tensor): Output tensor of the shape [B, C_out, H (// 2), W (// 2)].
      """
      # Save original shape
      _, c, h, w = data.shape
      # Perform partition
      input_partitioned = self.partition_function(
          data, self.grid_window_size
      )
      input_partitioned = input_partitioned.view(
          -1, self.grid_window_size[0] * self.grid_window_size[1], c
      )
      # Perform normalization, attention, and dropout
      output = input_partitioned + self.drop_path(
          self.attention(self.norm_1(input_partitioned))
      )
      # Perform normalization, MLP, and dropout
      output = output * self.drop_path(self.mlp(self.norm_2(output)))
      # Reverse partition
      output = self.reverse_function(
          output, (h, w), self.grid_window_size
      )
      return output
  
file_path: towhee/towhee/models/max_vit/max_vit_block.py
imports: from typing import Tuple, Union, List, Type, Callable
  import torch
  from torch import nn
  from timm.models.layers import Mlp, DropPath
  from towhee.models.layers.relative_self_attention import RelativeSelfAttention
  from towhee.models.layers.mbconv import MBConv
  from towhee.models.max_vit.max_vit_utils import window_reverse, window_partition, grid_reverse, grid_partition
label_desc: Use image node_id 2 for calling the MaxViTTransformerBlock obj's underlying member method code with example usage: obj.forward(data) and returns: output