class_docstr: The search space that is proposed in `AutoFormer <https://arxiv.org/abs/2107.00651>`__.
  There are four searchable variables: depth, embedding dimension, heads number and MLP ratio.
  
  Parameters
  ----------
  search_embed_dim
      The search space of embedding dimension. Use a list to specify search range.
  search_mlp_ratio
      The search space of MLP ratio. Use a list to specify search range.
  search_num_heads
      The search space of number of heads. Use a list to specify search range.
  search_depth
      The search space of depth. Use a list to specify search range.
  img_size
      Size of input image.
  patch_size
      Size of image patch.
  in_channels
      Number of channels of the input image.
  num_labels
      Number of classes for classifier.
  qkv_bias
      Whether to use bias item in the qkv embedding.
  drop_rate
      Drop rate of the MLP projection in MSA and FFN.
  attn_drop_rate
      Drop rate of attention.
  drop_path_rate
      Drop path rate.
  pre_norm
      Whether to use pre_norm. Otherwise post_norm is used.
  global_pooling
      Whether to use global pooling to generate the image representation. Otherwise the cls_token is used.
  absolute_position
      Whether to use absolute positional embeddings.
  qk_scale
      The scaler on score map in self-attention.
  rpe
      Whether to use relative position encoding.
code: def __init__(
      self,
      search_embed_dim: Tuple[int, ...] = (192, 216, 240),
      search_mlp_ratio: Tuple[float, ...] = (3.0, 3.5, 4.0),
      search_num_heads: Tuple[int, ...] = (3, 4),
      search_depth: Tuple[int, ...] = (12, 13, 14),
      img_size: int = 224,
      patch_size: int = 16,
      in_channels: int = 3,
      num_labels: int = 1000,
      qkv_bias: bool = True,
      drop_rate: float = 0.0,
      attn_drop_rate: float = 0.0,
      drop_path_rate: float = 0.1,
      pre_norm: bool = True,
      global_pooling: bool = True,
      absolute_position: bool = True,
      qk_scale: float | None = None,
      rpe: bool = True,
  ):
      super().__init__()
  
      # define search space parameters
      embed_dim = nni.choice("embed_dim", list(search_embed_dim))
      depth = nni.choice("depth", list(search_depth))
      mlp_ratios = [
          nni.choice(f"mlp_ratio_{i}", list(search_mlp_ratio))
          for i in range(max(search_depth))
      ]
      num_heads = [
          nni.choice(f"num_head_{i}", list(search_num_heads))
          for i in range(max(search_depth))
      ]
  
      self.patch_embed = MutableConv2d(
          in_channels,
          cast(int, embed_dim),
          kernel_size=patch_size,
          stride=patch_size,
      )
      self.patches_num = int((img_size // patch_size) ** 2)
      self.global_pooling = global_pooling
  
      self.cls_token = ClassToken(cast(int, embed_dim))
      self.pos_embed = (
          AbsolutePositionEmbedding(
              self.patches_num + 1, cast(int, embed_dim)
          )
          if absolute_position
          else nn.Identity()
      )
  
      dpr = [
          x.item()
          for x in torch.linspace(0, drop_path_rate, max(search_depth))
      ]  # stochastic depth decay rule
  
      self.blocks = Repeat(
          lambda index: TransformerEncoderLayer(
              embed_dim=embed_dim,
              num_heads=num_heads[index],
              mlp_ratio=mlp_ratios[index],
              qkv_bias=qkv_bias,
              drop_rate=drop_rate,
              attn_drop=attn_drop_rate,
              drop_path=dpr[index],
              rpe_length=img_size // patch_size,
              qk_scale=qk_scale,
              rpe=rpe,
              pre_norm=pre_norm,
          ),
          depth,
      )
  
      self.norm = (
          MutableLayerNorm(cast(int, embed_dim))
          if pre_norm
          else nn.Identity()
      )
      self.head = (
          MutableLinear(cast(int, embed_dim), num_labels)
          if num_labels > 0
          else nn.Identity()
      )
  
imports: from __future__ import annotations
  from typing import Tuple, cast, Any, Dict
  import torch
  import torch.nn
  import torch.nn.functional
  import nni
  from nni.mutable import Categorical, MutableExpression, frozen_context, ensure_frozen, Mutable
  from nni.nas.nn.pytorch import ParametrizedModule, ModelSpace, MutableModule, Repeat, MutableConv2d, MutableLinear, MutableLayerNorm
  from nni.nas.space import current_model
  from nni.nas.oneshot.pytorch.supermodule.operation import MixedOperation
  from nni.nas.profiler.pytorch.flops import FlopsResult
  from nni.nas.profiler.pytorch.utils import MutableShape, ShapeTensor, profiler_leaf_module
  from .utils.pretrained import load_pretrained_weight
  from .utils.nn import DropPath
label_desc: Use image node_id 1 to create a new AutoFormer object from inherited base classes: ModelSpace with example: obj = AutoFormer(search_embed_dim, search_mlp_ratio, search_num_heads, search_depth, img_size, patch_size, in_channels, num_labels, qkv_bias, drop_rate, attn_drop_rate, drop_path_rate, pre_norm, global_pooling, absolute_position, qk_scale, rpe)