code: def __init__(
      self,
      in_chs,
      out_chs,
      depth=1,
      downsample=True,
      attn_types=("spatial", "channel"),
      num_heads=3,
      window_size=7,
      mlp_ratio=4,
      qkv_bias=True,
      drop_path_rates=(0, 0),
      norm_layer=LayerNorm2d,
      norm_layer_cl=nn.LayerNorm,
      ffn=True,
      cpe_act=False,
  ):
      super().__init__()
  
      self.grad_checkpointing = False
  
      # downsample embedding layer at the beginning of each stage
      if downsample:
          self.downsample = Downsample(
              in_chs, out_chs, norm_layer=norm_layer
          )
      else:
          self.downsample = nn.Identity()
  
      """
           repeating alternating attention blocks in each stage
           default: (spatial -> channel) x depth
           
           potential opportunity to integrate with a more general version of ByobNet/ByoaNet
           since the logic is similar
          """
      stage_blocks = []
      for block_idx in range(depth):
          dual_attention_block = []
          for attn_idx, attn_type in enumerate(attn_types):
              if attn_type == "spatial":
                  dual_attention_block.append(
                      SpatialBlock(
                          dim=out_chs,
                          num_heads=num_heads,
                          mlp_ratio=mlp_ratio,
                          qkv_bias=qkv_bias,
                          drop_path=drop_path_rates[block_idx],
                          norm_layer=norm_layer_cl,
                          ffn=ffn,
                          cpe_act=cpe_act,
                          window_size=window_size,
                      )
                  )
              elif attn_type == "channel":
                  dual_attention_block.append(
                      ChannelBlock(
                          dim=out_chs,
                          num_heads=num_heads,
                          mlp_ratio=mlp_ratio,
                          qkv_bias=qkv_bias,
                          drop_path=drop_path_rates[block_idx],
                          norm_layer=norm_layer_cl,
                          ffn=ffn,
                          cpe_act=cpe_act,
                      )
                  )
          stage_blocks.append(nn.Sequential(*dual_attention_block))
      self.blocks = nn.Sequential(*stage_blocks)
  
imports: from functools import partial
  from typing import Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import DropPath, to_2tuple, trunc_normal_, Mlp, LayerNorm2d, get_norm_layer, use_fused_attn
  from timm.layers import NormMlpClassifierHead, ClassifierHead
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_function
  from ._manipulate import checkpoint_seq
  from ._registry import generate_default_cfgs, register_model
label_desc: Use image node_id 1 to create a new DaVitStage object from inherited base classes: nn with example: obj = DaVitStage(in_chs, out_chs, depth, downsample, attn_types, num_heads, window_size, mlp_ratio, qkv_bias, drop_path_rates, norm_layer, norm_layer_cl, ffn, cpe_act)