class_docstr: Implementation of metaformer block with MHSA as token mixer.
  For more details on Metaformer structure, please refer to:
  `MetaFormer Is Actually What You Need for Vision <https://arxiv.org/pdf/2111.11418.pdf>`_
  
  Args:
      opts: Command line arguments.
      dim: Number of embedding dimensions.
      mlp_ratio: MLP expansion ratio. Default: 4.0
      drop: Dropout rate. Default: 0.0
      drop_path: Drop path rate. Default: 0.0
      use_layer_scale: Flag to turn on layer scale. Default: ``True``
      layer_scale_init_value: Layer scale value at initialization. Default: 1e-5
class_docstr_tok: ["Implementation","of","metaformer","block","with","MHSA","as","token","mixer",".","For","more","details","on","Metaformer","structure",",","please","refer","to",":","`","MetaFormer","Is","Actually","What","You","Need","for","Vision","<","https",":","\/\/arxiv.org\/pdf\/2111.11418.pdf",">","`","_","Args",":","opts",":","Command","line","arguments",".","dim",":","Number","of","embedding","dimensions",".","mlp_ratio",":","MLP","expansion","ratio",".","Default",":","4.0","drop",":","Dropout","rate",".","Default",":","0.0","drop_path",":","Drop","path","rate",".","Default",":","0.0","use_layer_scale",":","Flag","to","turn","on","layer","scale",".","Default",":","``","True","``","layer_scale_init_value",":","Layer","scale","value","at","initialization",".","Default",":","1e-5"]
docstr: Perform appropriate reshaping before and after MHSA block.
  
  Args:
      x: Input tensor of shape :math:`(B, C, H, W)`.
  
  Returns:
      torch.Tensor of shape :math:`(B, C, H, W)`.
docstr_tok: ["Perform","appropriate","reshaping","before","and","after","MHSA","block",".","Args",":","x",":","Input","tensor","of","shape",":","math",":","`","(","B",",","C",",","H",",","W",")","`",".","Returns",":","torch.Tensor","of","shape",":","math",":","`","(","B",",","C",",","H",",","W",")","`","."]
code: def _apply_mhsa(self, x: torch.Tensor) -> torch.Tensor:
      """
      Perform appropriate reshaping before and after MHSA block.
  
      Args:
          x: Input tensor of shape :math:`(B, C, H, W)`.
  
      Returns:
          torch.Tensor of shape :math:`(B, C, H, W)`.
      """
      x_norm = self.norm(x)
      B, C, H, W = x_norm.shape
      x_norm_reshaped = torch.flatten(x_norm, start_dim=2).transpose(
          -2, -1
      )  # (B, N, C), where N = H * W
      out = self.token_mixer(x_norm_reshaped)
      out = out.transpose(-2, -1).reshape(B, C, H, W)
      return out
  
imports: import argparse
  from typing import Optional, Tuple, Union
  import torch
  import torch.nn
  from cvnets.layers import ConvLayer2d, MultiHeadAttention, StochasticDepth
  from cvnets.layers.normalization.batch_norm import BatchNorm2d
  from cvnets.modules import BaseModule
  from cvnets.modules.mobileone_block import MobileOneBlock, RepLKBlock
label_desc: Use image node_id 2 for calling the AttentionBlock obj's underlying member method code with example usage: obj._apply_mhsa(x) and returns: out