class_docstr: CoaT class. 
class_docstr_tok: ["CoaT","class","."]
code: def __init__(
      self,
      img_size=224,
      patch_size=16,
      in_chans=3,
      num_classes=1000,
      embed_dims=(64, 128, 320, 512),
      serial_depths=(3, 4, 6, 3),
      parallel_depth=0,
      num_heads=8,
      mlp_ratios=(4, 4, 4, 4),
      qkv_bias=True,
      drop_rate=0.0,
      proj_drop_rate=0.0,
      attn_drop_rate=0.0,
      drop_path_rate=0.0,
      norm_layer=LayerNorm,
      return_interm_layers=False,
      out_features=None,
      crpe_window=None,
      global_pool="token",
  ):
      super().__init__()
      assert global_pool in ("token", "avg")
      crpe_window = crpe_window or {3: 2, 5: 3, 7: 3}
      self.return_interm_layers = return_interm_layers
      self.out_features = out_features
      self.embed_dims = embed_dims
      self.num_features = embed_dims[-1]
      self.num_classes = num_classes
      self.global_pool = global_pool
  
      # Patch embeddings.
      img_size = to_2tuple(img_size)
      self.patch_embed1 = PatchEmbed(
          img_size=img_size,
          patch_size=patch_size,
          in_chans=in_chans,
          embed_dim=embed_dims[0],
          norm_layer=nn.LayerNorm,
      )
      self.patch_embed2 = PatchEmbed(
          img_size=[x // 4 for x in img_size],
          patch_size=2,
          in_chans=embed_dims[0],
          embed_dim=embed_dims[1],
          norm_layer=nn.LayerNorm,
      )
      self.patch_embed3 = PatchEmbed(
          img_size=[x // 8 for x in img_size],
          patch_size=2,
          in_chans=embed_dims[1],
          embed_dim=embed_dims[2],
          norm_layer=nn.LayerNorm,
      )
      self.patch_embed4 = PatchEmbed(
          img_size=[x // 16 for x in img_size],
          patch_size=2,
          in_chans=embed_dims[2],
          embed_dim=embed_dims[3],
          norm_layer=nn.LayerNorm,
      )
  
      # Class tokens.
      self.cls_token1 = nn.Parameter(torch.zeros(1, 1, embed_dims[0]))
      self.cls_token2 = nn.Parameter(torch.zeros(1, 1, embed_dims[1]))
      self.cls_token3 = nn.Parameter(torch.zeros(1, 1, embed_dims[2]))
      self.cls_token4 = nn.Parameter(torch.zeros(1, 1, embed_dims[3]))
  
      # Convolutional position encodings.
      self.cpe1 = ConvPosEnc(dim=embed_dims[0], k=3)
      self.cpe2 = ConvPosEnc(dim=embed_dims[1], k=3)
      self.cpe3 = ConvPosEnc(dim=embed_dims[2], k=3)
      self.cpe4 = ConvPosEnc(dim=embed_dims[3], k=3)
  
      # Convolutional relative position encodings.
      self.crpe1 = ConvRelPosEnc(
          head_chs=embed_dims[0] // num_heads,
          num_heads=num_heads,
          window=crpe_window,
      )
      self.crpe2 = ConvRelPosEnc(
          head_chs=embed_dims[1] // num_heads,
          num_heads=num_heads,
          window=crpe_window,
      )
      self.crpe3 = ConvRelPosEnc(
          head_chs=embed_dims[2] // num_heads,
          num_heads=num_heads,
          window=crpe_window,
      )
      self.crpe4 = ConvRelPosEnc(
          head_chs=embed_dims[3] // num_heads,
          num_heads=num_heads,
          window=crpe_window,
      )
  
      # Disable stochastic depth.
      dpr = drop_path_rate
      assert dpr == 0.0
      skwargs = dict(
          num_heads=num_heads,
          qkv_bias=qkv_bias,
          proj_drop=proj_drop_rate,
          attn_drop=attn_drop_rate,
          drop_path=dpr,
          norm_layer=norm_layer,
      )
  
      # Serial blocks 1.
      self.serial_blocks1 = nn.ModuleList(
          [
              SerialBlock(
                  dim=embed_dims[0],
                  mlp_ratio=mlp_ratios[0],
                  shared_cpe=self.cpe1,
                  shared_crpe=self.crpe1,
                  **skwargs,
              )
              for _ in range(serial_depths[0])
          ]
      )
  
      # Serial blocks 2.
      self.serial_blocks2 = nn.ModuleList(
          [
              SerialBlock(
                  dim=embed_dims[1],
                  mlp_ratio=mlp_ratios[1],
                  shared_cpe=self.cpe2,
                  shared_crpe=self.crpe2,
                  **skwargs,
              )
              for _ in range(serial_depths[1])
          ]
      )
  
      # Serial blocks 3.
      self.serial_blocks3 = nn.ModuleList(
          [
              SerialBlock(
                  dim=embed_dims[2],
                  mlp_ratio=mlp_ratios[2],
                  shared_cpe=self.cpe3,
                  shared_crpe=self.crpe3,
                  **skwargs,
              )
              for _ in range(serial_depths[2])
          ]
      )
  
      # Serial blocks 4.
      self.serial_blocks4 = nn.ModuleList(
          [
              SerialBlock(
                  dim=embed_dims[3],
                  mlp_ratio=mlp_ratios[3],
                  shared_cpe=self.cpe4,
                  shared_crpe=self.crpe4,
                  **skwargs,
              )
              for _ in range(serial_depths[3])
          ]
      )
  
      # Parallel blocks.
      self.parallel_depth = parallel_depth
      if self.parallel_depth > 0:
          self.parallel_blocks = nn.ModuleList(
              [
                  ParallelBlock(
                      dims=embed_dims,
                      mlp_ratios=mlp_ratios,
                      shared_crpes=(
                          self.crpe1,
                          self.crpe2,
                          self.crpe3,
                          self.crpe4,
                      ),
                      **skwargs,
                  )
                  for _ in range(parallel_depth)
              ]
          )
      else:
          self.parallel_blocks = None
  
      # Classification head(s).
      if not self.return_interm_layers:
          if self.parallel_blocks is not None:
              self.norm2 = norm_layer(embed_dims[1])
              self.norm3 = norm_layer(embed_dims[2])
          else:
              self.norm2 = self.norm3 = None
          self.norm4 = norm_layer(embed_dims[3])
  
          if self.parallel_depth > 0:
              # CoaT series: Aggregate features of last three scales for classification.
              assert embed_dims[1] == embed_dims[2] == embed_dims[3]
              self.aggregate = torch.nn.Conv1d(
                  in_channels=3, out_channels=1, kernel_size=1
              )
              self.head_drop = nn.Dropout(drop_rate)
              self.head = (
                  nn.Linear(self.num_features, num_classes)
                  if num_classes > 0
                  else nn.Identity()
              )
          else:
              # CoaT-Lite series: Use feature of last scale for classification.
              self.aggregate = None
              self.head_drop = nn.Dropout(drop_rate)
              self.head = (
                  nn.Linear(self.num_features, num_classes)
                  if num_classes > 0
                  else nn.Identity()
              )
  
      # Initialize weights.
      trunc_normal_(self.cls_token1, std=0.02)
      trunc_normal_(self.cls_token2, std=0.02)
      trunc_normal_(self.cls_token3, std=0.02)
      trunc_normal_(self.cls_token4, std=0.02)
      self.apply(self._init_weights)
  
imports: from functools import partial
  from typing import Tuple, List, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import PatchEmbed, Mlp, DropPath, to_2tuple, trunc_normal_, _assert, LayerNorm
  from ._builder import build_model_with_cfg
  from ._registry import register_model, generate_default_cfgs
label_desc: Use image node_id 1 to create a new CoaT object from inherited base classes: nn with example: obj = CoaT(img_size, patch_size, in_chans, num_classes, embed_dims, serial_depths, parallel_depth, num_heads, mlp_ratios, qkv_bias, drop_rate, proj_drop_rate, attn_drop_rate, drop_path_rate, norm_layer, return_interm_layers, out_features, crpe_window, global_pool)