class_docstr: Vision Transformer
  A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
      - https://arxiv.org/abs/2010.11929
class_docstr_tok: ["Vision","Transformer","A","PyTorch","impl","of",":","`","An","Image","is","Worth","16x16","Words",":","Transformers","for","Image","Recognition","at","Scale","`","-","https",":","\/\/arxiv.org\/abs\/2010.11929"]
docstr: Args:
      img_size (int, tuple): input image size
      patch_size (int, tuple): patch size
      in_chans (int): number of input channels
      num_classes (int): number of classes for classification head
      global_pool (str): type of global pooling for final sequence (default: 'token')
      embed_dim (int): embedding dimension
      depth (int): depth of transformer
      num_heads (int): number of attention heads
      mlp_ratio (int): ratio of mlp hidden dim to embedding dim
      qkv_bias (bool): enable bias for qkv if True
      init_values: (float): layer-scale init values
      class_token (bool): use class token
      fc_norm (Optional[bool]): pre-fc norm after pool, set if global_pool == 'avg' if None (default: None)
      drop_rate (float): dropout rate
      attn_drop_rate (float): attention dropout rate
      drop_path_rate (float): stochastic depth rate
      weight_init (str): weight init scheme
      embed_layer (nn.Module): patch embedding layer
      norm_layer: (nn.Module): normalization layer
      act_layer: (nn.Module): MLP activation layer
docstr_tok: ["Args",":","img_size","(","int",",","tuple",")",":","input","image","size","patch_size","(","int",",","tuple",")",":","patch","size","in_chans","(","int",")",":","number","of","input","channels","num_classes","(","int",")",":","number","of","classes","for","classification","head","global_pool","(","str",")",":","type","of","global","pooling","for","final","sequence","(","default",":","'token","'",")","embed_dim","(","int",")",":","embedding","dimension","depth","(","int",")",":","depth","of","transformer","num_heads","(","int",")",":","number","of","attention","heads","mlp_ratio","(","int",")",":","ratio","of","mlp","hidden","dim","to","embedding","dim","qkv_bias","(","bool",")",":","enable","bias","for","qkv","if","True","init_values",":","(","float",")",":","layer-scale","init","values","class_token","(","bool",")",":","use","class","token","fc_norm","(","Optional","[","bool","]",")",":","pre-fc","norm","after","pool",",","set","if","global_pool","==","'avg","'","if","None","(","default",":","None",")","drop_rate","(","float",")",":","dropout","rate","attn_drop_rate","(","float",")",":","attention","dropout","rate","drop_path_rate","(","float",")",":","stochastic","depth","rate","weight_init","(","str",")",":","weight","init","scheme","embed_layer","(","nn.Module",")",":","patch","embedding","layer","norm_layer",":","(","nn.Module",")",":","normalization","layer","act_layer",":","(","nn.Module",")",":","MLP","activation","layer"]
code: def __init__(
      self,
      img_size=224,
      patch_size=16,
      in_chans=3,
      num_classes=1000,
      global_pool="token",
      embed_dim=768,
      depth=12,
      num_heads=12,
      mlp_ratio=4.0,
      qkv_bias=True,
      init_values=None,
      class_token=True,
      no_embed_class=False,
      pre_norm=False,
      fc_norm=None,
      drop_rate=0.0,
      attn_drop_rate=0.0,
      drop_path_rate=0.0,
      weight_init="",
      embed_layer=PatchEmbed,
      norm_layer=None,
      act_layer=None,
      use_flash_attn=False,
      fused_bias_fc=False,
      fused_dense_gelu_dense=False,
      fused_dropout_add_ln=False,
      bf16=False,
  ):
      """
      Args:
          img_size (int, tuple): input image size
          patch_size (int, tuple): patch size
          in_chans (int): number of input channels
          num_classes (int): number of classes for classification head
          global_pool (str): type of global pooling for final sequence (default: 'token')
          embed_dim (int): embedding dimension
          depth (int): depth of transformer
          num_heads (int): number of attention heads
          mlp_ratio (int): ratio of mlp hidden dim to embedding dim
          qkv_bias (bool): enable bias for qkv if True
          init_values: (float): layer-scale init values
          class_token (bool): use class token
          fc_norm (Optional[bool]): pre-fc norm after pool, set if global_pool == 'avg' if None (default: None)
          drop_rate (float): dropout rate
          attn_drop_rate (float): attention dropout rate
          drop_path_rate (float): stochastic depth rate
          weight_init (str): weight init scheme
          embed_layer (nn.Module): patch embedding layer
          norm_layer: (nn.Module): normalization layer
          act_layer: (nn.Module): MLP activation layer
      """
      super().__init__()
      assert (
          global_pool == "token"
      ), "Only support pooling with CLS token"
      assert class_token
      assert init_values is None, "LayerScale is not supported yet"
      assert weight_init == ""
      assert fc_norm is None
      # pre_norm seems redundant, as there's a LayerNorm right at the start of each block, idk
      assert not pre_norm
      use_fc_norm = global_pool == "avg" if fc_norm is None else fc_norm
      norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
      act_layer = act_layer or nn.GELU
  
      self.num_classes = num_classes
      self.global_pool = global_pool
      self.num_features = (
          self.embed_dim
      ) = embed_dim  # num_features for consistency with other models
      self.num_prefix_tokens = 1 if class_token else 0
      self.no_embed_class = no_embed_class
  
      self.patch_embed = embed_layer(
          img_size=img_size,
          patch_size=patch_size,
          in_chans=in_chans,
          embed_dim=embed_dim,
          bias=not pre_norm,  # disable bias if pre-norm is used (e.g. CLIP)
      )
      num_patches = self.patch_embed.num_patches
  
      self.cls_token = (
          nn.Parameter(torch.zeros(1, 1, embed_dim))
          if class_token
          else None
      )
      embed_len = (
          num_patches
          if no_embed_class
          else num_patches + self.num_prefix_tokens
      )
      self.pos_embed = nn.Parameter(
          torch.randn(1, embed_len, embed_dim) * 0.02
      )
      self.pos_drop = nn.Dropout(p=drop_rate)
  
      dpr = [
          x.item() for x in torch.linspace(0, drop_path_rate, depth)
      ]  # stochastic depth decay rule
  
      # We change the order of residual and layer norm:
      # Instead of LN -> Attn / MLP -> Dropout -> Add, we do:
      # Attn / MLP -> Dropout -> Add -> LN, returning both the residual branch (output of Add) and
      # the main branch (output of LN). The model definition is unchanged, but the mapping of the
      # nn.LayerNorm weights are changed.
      # This is for performance reason: we can fuse dropout + add + layer_norm.
      # self.norm_0 is the first layer norm in the model, while self.norm
      # (in the pretrained weight) is the final layer norm.
      self.norm_0 = norm_layer(embed_dim)
  
      self.blocks = nn.ModuleList(
          [
              create_block(
                  embed_dim,
                  num_heads,
                  mlp_ratio,
                  qkv_bias,
                  drop_rate,
                  attn_drop_rate,
                  drop_path=dpr[i],
                  norm_layer=norm_layer,
                  act_layer=act_layer,
                  use_flash_attn=use_flash_attn,
                  fused_bias_fc=fused_bias_fc,
                  fused_dense_gelu_dense=fused_dense_gelu_dense,
                  fused_dropout_add_ln=fused_dropout_add_ln,
                  bf16=bf16,
              )
              for i in range(depth)
          ]
      )
  
      # Classifier Head
      self.head = (
          nn.Linear(self.embed_dim, num_classes)
          if num_classes > 0
          else nn.Identity()
      )
  
      self.init_weights(weight_init)
  
imports: import math
  from functools import partial
  from copy import deepcopy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.init import trunc_normal_
  from timm.models.helpers import named_apply
  from timm.models.layers import PatchEmbed
  from src.models.modules.mha import MHA
  from src.models.modules.mlp import Mlp, FusedDenseGeluDense
  from src.models.modules.block import Block
label_desc: Use image node_id 1 to create a new VisionTransformer object from inherited base classes: nn with example: obj = VisionTransformer(img_size, patch_size, in_chans, num_classes, global_pool, embed_dim, depth, num_heads, mlp_ratio, qkv_bias, init_values, class_token, no_embed_class, pre_norm, fc_norm, drop_rate, attn_drop_rate, drop_path_rate, weight_init, embed_layer, norm_layer, act_layer, use_flash_attn, fused_bias_fc, fused_dense_gelu_dense, fused_dropout_add_ln, bf16)