code: def __init__(
      self,
      in_chans=3,
      num_classes=1000,
      global_pool="avg",
      dims=(24, 48, 88, 168),
      depths=(3, 3, 9, 3),
      global_block_counts=(0, 1, 1, 1),
      kernel_sizes=(3, 5, 7, 9),
      heads=(8, 8, 8, 8),
      d2_scales=(2, 2, 3, 4),
      use_pos_emb=(False, True, False, False),
      ls_init_value=1e-6,
      head_init_scale=1.0,
      expand_ratio=4,
      downsample_block=False,
      conv_bias=True,
      stem_type="patch",
      head_norm_first=False,
      act_layer=nn.GELU,
      drop_path_rate=0.0,
      drop_rate=0.0,
  ):
      super().__init__()
      self.num_classes = num_classes
      self.global_pool = global_pool
      self.drop_rate = drop_rate
      norm_layer = partial(LayerNorm2d, eps=1e-6)
      norm_layer_cl = partial(nn.LayerNorm, eps=1e-6)
      self.feature_info = []
  
      assert stem_type in ("patch", "overlap")
      if stem_type == "patch":
          self.stem = nn.Sequential(
              nn.Conv2d(
                  in_chans,
                  dims[0],
                  kernel_size=4,
                  stride=4,
                  bias=conv_bias,
              ),
              norm_layer(dims[0]),
          )
      else:
          self.stem = nn.Sequential(
              nn.Conv2d(
                  in_chans,
                  dims[0],
                  kernel_size=9,
                  stride=4,
                  padding=9 // 2,
                  bias=conv_bias,
              ),
              norm_layer(dims[0]),
          )
  
      curr_stride = 4
      stages = []
      dp_rates = [
          x.tolist()
          for x in torch.linspace(0, drop_path_rate, sum(depths)).split(
              depths
          )
      ]
      in_chs = dims[0]
      for i in range(4):
          stride = 2 if curr_stride == 2 or i > 0 else 1
          # FIXME support dilation / output_stride
          curr_stride *= stride
          stages.append(
              EdgeNeXtStage(
                  in_chs=in_chs,
                  out_chs=dims[i],
                  stride=stride,
                  depth=depths[i],
                  num_global_blocks=global_block_counts[i],
                  num_heads=heads[i],
                  drop_path_rates=dp_rates[i],
                  scales=d2_scales[i],
                  expand_ratio=expand_ratio,
                  kernel_size=kernel_sizes[i],
                  use_pos_emb=use_pos_emb[i],
                  ls_init_value=ls_init_value,
                  downsample_block=downsample_block,
                  conv_bias=conv_bias,
                  norm_layer=norm_layer,
                  norm_layer_cl=norm_layer_cl,
                  act_layer=act_layer,
              )
          )
          # NOTE feature_info use currently assumes stage 0 == stride 1, rest are stride 2
          in_chs = dims[i]
          self.feature_info += [
              dict(
                  num_chs=in_chs,
                  reduction=curr_stride,
                  module=f"stages.{i}",
              )
          ]
  
      self.stages = nn.Sequential(*stages)
  
      self.num_features = dims[-1]
      if head_norm_first:
          self.norm_pre = norm_layer(self.num_features)
          self.head = ClassifierHead(
              self.num_features,
              num_classes,
              pool_type=global_pool,
              drop_rate=self.drop_rate,
          )
      else:
          self.norm_pre = nn.Identity()
          self.head = NormMlpClassifierHead(
              self.num_features,
              num_classes,
              pool_type=global_pool,
              drop_rate=self.drop_rate,
              norm_layer=norm_layer,
          )
  
      named_apply(
          partial(_init_weights, head_init_scale=head_init_scale), self
      )
  
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 EdgeNeXt object from inherited base classes: nn with example: obj = EdgeNeXt(in_chans, num_classes, global_pool, dims, depths, global_block_counts, kernel_sizes, heads, d2_scales, use_pos_emb, ls_init_value, head_init_scale, expand_ratio, downsample_block, conv_bias, stem_type, head_norm_first, act_layer, drop_path_rate, drop_rate)