class_docstr: Nested Transformer (NesT)
  
  A PyTorch impl of : `Aggregating Nested Transformers`
      - https://arxiv.org/abs/2105.12723
class_docstr_tok: ["Nested","Transformer","(","NesT",")","A","PyTorch","impl","of",":","`","Aggregating","Nested","Transformers","`","-","https",":","\/\/arxiv.org\/abs\/2105.12723"]
docstr: Args:
      img_size (int, tuple): input image size
      in_chans (int): number of input channels
      patch_size (int): patch size
      num_levels (int): number of block hierarchies (T_d in the paper)
      embed_dims (int, tuple): embedding dimensions of each level
      num_heads (int, tuple): number of attention heads for each level
      depths (int, tuple): number of transformer layers for each level
      num_classes (int): number of classes for classification head
      mlp_ratio (int): ratio of mlp hidden dim to embedding dim for MLP of transformer layers
      qkv_bias (bool): enable bias for qkv if True
      drop_rate (float): dropout rate for MLP of transformer layers, MSA final projection layer, and classifier
      attn_drop_rate (float): attention dropout rate
      drop_path_rate (float): stochastic depth rate
      norm_layer: (nn.Module): normalization layer for transformer layers
      act_layer: (nn.Module): activation layer in MLP of transformer layers
      pad_type: str: Type of padding to use '' for PyTorch symmetric, 'same' for TF SAME
      weight_init: (str): weight init scheme
      global_pool: (str): type of pooling operation to apply to final feature map
  
  Notes:
      - Default values follow NesT-B from the original Jax code.
      - `embed_dims`, `num_heads`, `depths` should be ints or tuples with length `num_levels`.
      - For those following the paper, Table A1 may have errors!
          - https://github.com/google-research/nested-transformer/issues/2
code: def __init__(
      self,
      img_size=224,
      in_chans=3,
      patch_size=4,
      num_levels=3,
      embed_dims=(128, 256, 512),
      num_heads=(4, 8, 16),
      depths=(2, 2, 20),
      num_classes=1000,
      mlp_ratio=4.0,
      qkv_bias=True,
      drop_rate=0.0,
      proj_drop_rate=0.0,
      attn_drop_rate=0.0,
      drop_path_rate=0.5,
      norm_layer=None,
      act_layer=None,
      pad_type="",
      weight_init="",
      global_pool="avg",
  ):
      """
      Args:
          img_size (int, tuple): input image size
          in_chans (int): number of input channels
          patch_size (int): patch size
          num_levels (int): number of block hierarchies (T_d in the paper)
          embed_dims (int, tuple): embedding dimensions of each level
          num_heads (int, tuple): number of attention heads for each level
          depths (int, tuple): number of transformer layers for each level
          num_classes (int): number of classes for classification head
          mlp_ratio (int): ratio of mlp hidden dim to embedding dim for MLP of transformer layers
          qkv_bias (bool): enable bias for qkv if True
          drop_rate (float): dropout rate for MLP of transformer layers, MSA final projection layer, and classifier
          attn_drop_rate (float): attention dropout rate
          drop_path_rate (float): stochastic depth rate
          norm_layer: (nn.Module): normalization layer for transformer layers
          act_layer: (nn.Module): activation layer in MLP of transformer layers
          pad_type: str: Type of padding to use '' for PyTorch symmetric, 'same' for TF SAME
          weight_init: (str): weight init scheme
          global_pool: (str): type of pooling operation to apply to final feature map
  
      Notes:
          - Default values follow NesT-B from the original Jax code.
          - `embed_dims`, `num_heads`, `depths` should be ints or tuples with length `num_levels`.
          - For those following the paper, Table A1 may have errors!
              - https://github.com/google-research/nested-transformer/issues/2
      """
      super().__init__()
  
      for param_name in ["embed_dims", "num_heads", "depths"]:
          param_value = locals()[param_name]
          if isinstance(param_value, collections.abc.Sequence):
              assert (
                  len(param_value) == num_levels
              ), f"Require `len({param_name}) == num_levels`"
  
      embed_dims = to_ntuple(num_levels)(embed_dims)
      num_heads = to_ntuple(num_levels)(num_heads)
      depths = to_ntuple(num_levels)(depths)
      self.num_classes = num_classes
      self.num_features = embed_dims[-1]
      self.feature_info = []
      norm_layer = norm_layer or LayerNorm
      act_layer = act_layer or nn.GELU
      self.drop_rate = drop_rate
      self.num_levels = num_levels
      if isinstance(img_size, collections.abc.Sequence):
          assert (
              img_size[0] == img_size[1]
          ), "Model only handles square inputs"
          img_size = img_size[0]
      assert (
          img_size % patch_size == 0
      ), "`patch_size` must divide `img_size` evenly"
      self.patch_size = patch_size
  
      # Number of blocks at each level
      self.num_blocks = (4 ** torch.arange(num_levels)).flip(0).tolist()
      assert (img_size // patch_size) % math.sqrt(
          self.num_blocks[0]
      ) == 0, "First level blocks don't fit evenly. Check `img_size`, `patch_size`, and `num_levels`"
  
      # Block edge size in units of patches
      # Hint: (img_size // patch_size) gives number of patches along edge of image. sqrt(self.num_blocks[0]) is the
      #  number of blocks along edge of image
      self.block_size = int(
          (img_size // patch_size) // math.sqrt(self.num_blocks[0])
      )
  
      # Patch embedding
      self.patch_embed = PatchEmbed(
          img_size=img_size,
          patch_size=patch_size,
          in_chans=in_chans,
          embed_dim=embed_dims[0],
          flatten=False,
      )
      self.num_patches = self.patch_embed.num_patches
      self.seq_length = self.num_patches // self.num_blocks[0]
  
      # Build up each hierarchical level
      levels = []
      dp_rates = [
          x.tolist()
          for x in torch.linspace(0, drop_path_rate, sum(depths)).split(
              depths
          )
      ]
      prev_dim = None
      curr_stride = 4
      for i in range(len(self.num_blocks)):
          dim = embed_dims[i]
          levels.append(
              NestLevel(
                  self.num_blocks[i],
                  self.block_size,
                  self.seq_length,
                  num_heads[i],
                  depths[i],
                  dim,
                  prev_dim,
                  mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias,
                  proj_drop=proj_drop_rate,
                  attn_drop=attn_drop_rate,
                  drop_path=dp_rates[i],
                  norm_layer=norm_layer,
                  act_layer=act_layer,
                  pad_type=pad_type,
              )
          )
          self.feature_info += [
              dict(
                  num_chs=dim,
                  reduction=curr_stride,
                  module=f"levels.{i}",
              )
          ]
          prev_dim = dim
          curr_stride *= 2
      self.levels = nn.Sequential(*levels)
  
      # Final normalization layer
      self.norm = norm_layer(embed_dims[-1])
  
      # Classifier
      global_pool, head = create_classifier(
          self.num_features, self.num_classes, pool_type=global_pool
      )
      self.global_pool = global_pool
      self.head_drop = nn.Dropout(drop_rate)
      self.head = head
  
      self.init_weights(weight_init)
  
imports: import collections.abc
  import logging
  import math
  from functools import partial
  import torch
  import torch.nn.functional
  from torch import nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import PatchEmbed, Mlp, DropPath, create_classifier, trunc_normal_, _assert
  from timm.layers import create_conv2d, create_pool2d, to_ntuple, use_fused_attn, LayerNorm
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_function
  from ._manipulate import checkpoint_seq, named_apply
  from ._registry import register_model, generate_default_cfgs, register_model_deprecations
label_desc: Use image node_id 1 to create a new Nest object from inherited base classes: nn with example: obj = Nest(img_size, in_chans, patch_size, num_levels, embed_dims, num_heads, depths, num_classes, mlp_ratio, qkv_bias, drop_rate, proj_drop_rate, attn_drop_rate, drop_path_rate, norm_layer, act_layer, pad_type, weight_init, global_pool)