code: def __init__(
      self,
      in_chs,
      out_chs,
      stride=2,
      depth=2,
      num_global_blocks=1,
      num_heads=4,
      scales=2,
      kernel_size=7,
      expand_ratio=4,
      use_pos_emb=False,
      downsample_block=False,
      conv_bias=True,
      ls_init_value=1.0,
      drop_path_rates=None,
      norm_layer=LayerNorm2d,
      norm_layer_cl=partial(nn.LayerNorm, eps=1e-6),
      act_layer=nn.GELU,
  ):
      super().__init__()
      self.grad_checkpointing = False
  
      if downsample_block or stride == 1:
          self.downsample = nn.Identity()
      else:
          self.downsample = nn.Sequential(
              norm_layer(in_chs),
              nn.Conv2d(
                  in_chs,
                  out_chs,
                  kernel_size=2,
                  stride=2,
                  bias=conv_bias,
              ),
          )
          in_chs = out_chs
  
      stage_blocks = []
      for i in range(depth):
          if i < depth - num_global_blocks:
              stage_blocks.append(
                  ConvBlock(
                      dim=in_chs,
                      dim_out=out_chs,
                      stride=stride
                      if downsample_block and i == 0
                      else 1,
                      conv_bias=conv_bias,
                      kernel_size=kernel_size,
                      expand_ratio=expand_ratio,
                      ls_init_value=ls_init_value,
                      drop_path=drop_path_rates[i],
                      norm_layer=norm_layer_cl,
                      act_layer=act_layer,
                  )
              )
          else:
              stage_blocks.append(
                  SplitTransposeBlock(
                      dim=in_chs,
                      num_scales=scales,
                      num_heads=num_heads,
                      expand_ratio=expand_ratio,
                      use_pos_emb=use_pos_emb,
                      conv_bias=conv_bias,
                      ls_init_value=ls_init_value,
                      drop_path=drop_path_rates[i],
                      norm_layer=norm_layer_cl,
                      act_layer=act_layer,
                  )
              )
          in_chs = out_chs
      self.blocks = nn.Sequential(*stage_blocks)
  
imports: import math
  from collections import OrderedDict
  from functools import partial
  from typing import Tuple
  import torch
  import torch.nn.functional
  from torch import nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import trunc_normal_tf_, DropPath, LayerNorm2d, Mlp, SelectAdaptivePool2d, create_conv2d, use_fused_attn, NormMlpClassifierHead, ClassifierHead
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_module
  from ._manipulate import named_apply, checkpoint_seq
  from ._registry import register_model, generate_default_cfgs
label_desc: Use image node_id 1 to create a new EdgeNeXtStage object from inherited base classes: nn with example: obj = EdgeNeXtStage(in_chs, out_chs, stride, depth, num_global_blocks, num_heads, scales, kernel_size, expand_ratio, use_pos_emb, downsample_block, conv_bias, ls_init_value, drop_path_rates, norm_layer, norm_layer_cl, act_layer)