class_docstr: Vision Transformer for Segment-Anything Model(SAM)
  
  A PyTorch impl of : `Exploring Plain Vision Transformer Backbones for Object Detection` or `Segment Anything Model (SAM)`
      - https://arxiv.org/abs/2010.11929
class_docstr_tok: ["Vision","Transformer","for","Segment-Anything","Model","(","SAM",")","A","PyTorch","impl","of",":","`","Exploring","Plain","Vision","Transformer","Backbones","for","Object","Detection","`","or","`","Segment","Anything","Model","(","SAM",")","`","-","https",":","\/\/arxiv.org\/abs\/2010.11929"]
docstr: Args:
      img_size: Input image size.
      patch_size: Patch size.
      in_chans: Number of image input channels.
      num_classes: Mumber of classes for classification head.
      global_pool: Type of global pooling for final sequence (default: 'token').
      embed_dim: Transformer embedding dimension.
      depth: Depth of transformer.
      num_heads: Number of attention heads.
      mlp_ratio: Ratio of mlp hidden dim to embedding dim.
      qkv_bias: Enable bias for qkv projections if True.
      init_values: Layer-scale init values (layer-scale enabled if not None).
      drop_rate: Head dropout rate.
      pos_drop_rate: Position embedding dropout rate.
      attn_drop_rate: Attention dropout rate.
      drop_path_rate: Stochastic depth rate.
      weight_init: Weight initialization scheme.
      embed_layer: Patch embedding layer.
      norm_layer: Normalization layer.
      act_layer: MLP activation layer.
      block_fn: Transformer block layer.
      use_abs_pos: If True, use absolute positional embeddings.
      use_rel_pos: If True, add relative positional embeddings to the attention map.
      use_rope: If True, add rotary position embeddings to q/k in attention block.
      window_size: Window size for window attention blocks. If 0, not use window attention.
      global_attn_indexes: Indexes for blocks using global attention. Used when window_size > 0.
      global_pool: Global pooling type.
      head_hidden_size: If set, use NormMlpHead
      ref_feat_shape: Tuple of reference feature shapes for ROPE, (global, local)
docstr_tok: ["Args",":","img_size",":","Input","image","size",".","patch_size",":","Patch","size",".","in_chans",":","Number","of","image","input","channels",".","num_classes",":","Mumber","of","classes","for","classification","head",".","global_pool",":","Type","of","global","pooling","for","final","sequence","(","default",":","'token","'",")",".","embed_dim",":","Transformer","embedding","dimension",".","depth",":","Depth","of","transformer",".","num_heads",":","Number","of","attention","heads",".","mlp_ratio",":","Ratio","of","mlp","hidden","dim","to","embedding","dim",".","qkv_bias",":","Enable","bias","for","qkv","projections","if","True",".","init_values",":","Layer-scale","init","values","(","layer-scale","enabled","if","not","None",")",".","drop_rate",":","Head","dropout","rate",".","pos_drop_rate",":","Position","embedding","dropout","rate",".","attn_drop_rate",":","Attention","dropout","rate",".","drop_path_rate",":","Stochastic","depth","rate",".","weight_init",":","Weight","initialization","scheme",".","embed_layer",":","Patch","embedding","layer",".","norm_layer",":","Normalization","layer",".","act_layer",":","MLP","activation","layer",".","block_fn",":","Transformer","block","layer",".","use_abs_pos",":","If","True",",","use","absolute","positional","embeddings",".","use_rel_pos",":","If","True",",","add","relative","positional","embeddings","to","the","attention","map",".","use_rope",":","If","True",",","add","rotary","position","embeddings","to","q\/k","in","attention","block",".","window_size",":","Window","size","for","window","attention","blocks",".","If","0",",","not","use","window","attention",".","global_attn_indexes",":","Indexes","for","blocks","using","global","attention",".","Used","when","window_size",">","0.","global_pool",":","Global","pooling","type",".","head_hidden_size",":","If","set",",","use","NormMlpHead","ref_feat_shape",":","Tuple","of","reference","feature","shapes","for","ROPE",",","(","global",",","local",")"]
code: def __init__(
      self,
      img_size: int = 1024,
      patch_size: int = 16,
      in_chans: int = 3,
      num_classes: int = 768,
      embed_dim: int = 768,
      depth: int = 12,
      num_heads: int = 12,
      mlp_ratio: float = 4.0,
      qkv_bias: bool = True,
      qk_norm: bool = False,
      init_values: Optional[float] = None,
      pre_norm: bool = False,
      drop_rate: float = 0.0,
      pos_drop_rate: float = 0.0,
      patch_drop_rate: float = 0.0,
      proj_drop_rate: float = 0.0,
      attn_drop_rate: float = 0.0,
      drop_path_rate: float = 0.0,
      weight_init: str = "",
      embed_layer: Callable = partial(
          PatchEmbed, output_fmt=Format.NHWC, strict_img_size=False
      ),
      norm_layer: Optional[Callable] = nn.LayerNorm,
      act_layer: Optional[Callable] = nn.GELU,
      block_fn: Callable = Block,
      mlp_layer: Callable = Mlp,
      use_abs_pos: bool = True,
      use_rel_pos: bool = False,
      use_rope: bool = False,
      window_size: int = 14,
      global_attn_indexes: Tuple[int, ...] = (),
      neck_chans: int = 256,
      global_pool: str = "avg",
      head_hidden_size: Optional[int] = None,
      ref_feat_shape: Optional[
          Tuple[Tuple[int, int], Tuple[int, int]]
      ] = None,
  ):
      """
      Args:
          img_size: Input image size.
          patch_size: Patch size.
          in_chans: Number of image input channels.
          num_classes: Mumber of classes for classification head.
          global_pool: Type of global pooling for final sequence (default: 'token').
          embed_dim: Transformer embedding dimension.
          depth: Depth of transformer.
          num_heads: Number of attention heads.
          mlp_ratio: Ratio of mlp hidden dim to embedding dim.
          qkv_bias: Enable bias for qkv projections if True.
          init_values: Layer-scale init values (layer-scale enabled if not None).
          drop_rate: Head dropout rate.
          pos_drop_rate: Position embedding dropout rate.
          attn_drop_rate: Attention dropout rate.
          drop_path_rate: Stochastic depth rate.
          weight_init: Weight initialization scheme.
          embed_layer: Patch embedding layer.
          norm_layer: Normalization layer.
          act_layer: MLP activation layer.
          block_fn: Transformer block layer.
          use_abs_pos: If True, use absolute positional embeddings.
          use_rel_pos: If True, add relative positional embeddings to the attention map.
          use_rope: If True, add rotary position embeddings to q/k in attention block.
          window_size: Window size for window attention blocks. If 0, not use window attention.
          global_attn_indexes: Indexes for blocks using global attention. Used when window_size > 0.
          global_pool: Global pooling type.
          head_hidden_size: If set, use NormMlpHead
          ref_feat_shape: Tuple of reference feature shapes for ROPE, (global, local)
      """
      super().__init__()
      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
      # num_features for consistency with other models
      self.num_features = self.embed_dim = embed_dim
      self.grad_checkpointing = False
  
      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
      )
      grid_size = self.patch_embed.grid_size
      if use_abs_pos:
          # Initialize absolute positional embedding with pretrain image size.
          self.pos_embed = nn.Parameter(
              torch.zeros(1, grid_size[0], grid_size[1], embed_dim)
          )
      else:
          self.pos_embed = None
      self.pos_drop = nn.Dropout(p=pos_drop_rate)
      if patch_drop_rate > 0:
          self.patch_drop = PatchDropout(
              patch_drop_rate,
              num_prefix_tokens=0,
          )
      else:
          self.patch_drop = nn.Identity()
      self.norm_pre = (
          norm_layer(embed_dim) if pre_norm else nn.Identity()
      )
  
      if use_rope:
          assert (
              not use_rel_pos
          ), "ROPE and relative pos embeddings should not be enabled at same time"
          if ref_feat_shape is not None:
              assert len(ref_feat_shape) == 2
              ref_feat_shape_global = to_2tuple(ref_feat_shape[0])
              ref_feat_shape_window = to_2tuple(ref_feat_shape[1])
          else:
              ref_feat_shape_global = ref_feat_shape_window = None
          self.rope_global = RotaryEmbeddingCat(
              embed_dim // num_heads,
              in_pixels=False,
              feat_shape=grid_size,
              ref_feat_shape=ref_feat_shape_global,
          )
          self.rope_window = RotaryEmbeddingCat(
              embed_dim // num_heads,
              in_pixels=False,
              feat_shape=to_2tuple(window_size),
              ref_feat_shape=ref_feat_shape_window,
          )
      else:
          self.rope_global = None
          self.rope_window = None
  
      # stochastic depth decay rule
      dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
      self.blocks = nn.Sequential(
          *[
              block_fn(
                  dim=embed_dim,
                  num_heads=num_heads,
                  mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias,
                  qk_norm=qk_norm,
                  init_values=init_values,
                  proj_drop=proj_drop_rate,
                  attn_drop=attn_drop_rate,
                  drop_path=dpr[i],
                  norm_layer=norm_layer,
                  act_layer=act_layer,
                  mlp_layer=mlp_layer,
                  use_rel_pos=use_rel_pos,
                  window_size=window_size
                  if i not in global_attn_indexes
                  else 0,
                  input_size=grid_size,
                  rope=self.rope_window
                  if i not in global_attn_indexes
                  else self.rope_global,
              )
              for i in range(depth)
          ]
      )
  
      if neck_chans:
          self.neck = nn.Sequential(
              nn.Conv2d(
                  embed_dim,
                  neck_chans,
                  kernel_size=1,
                  bias=False,
              ),
              LayerNorm2d(neck_chans),
              nn.Conv2d(
                  neck_chans,
                  neck_chans,
                  kernel_size=3,
                  padding=1,
                  bias=False,
              ),
              LayerNorm2d(neck_chans),
          )
          self.num_features = neck_chans
      else:
          if head_hidden_size:
              self.neck = nn.Identity()
          else:
              # should have a final norm with standard ClassifierHead
              self.neck = LayerNorm2d(embed_dim)
          neck_chans = embed_dim
  
      # Classifier Head
      if head_hidden_size:
          self.head = NormMlpClassifierHead(
              neck_chans,
              num_classes,
              hidden_size=head_hidden_size,
              pool_type=global_pool,
              drop_rate=drop_rate,
          )
      else:
          self.head = ClassifierHead(
              neck_chans,
              num_classes,
              pool_type=global_pool,
              drop_rate=drop_rate,
          )
  
imports: import logging
  from functools import partial
  from typing import Callable, Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch.jit import Final
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
  from timm.layers import PatchEmbed, Mlp, DropPath, PatchDropout, LayerNorm2d, ClassifierHead, NormMlpClassifierHead, Format, resample_abs_pos_embed_nhwc, RotaryEmbeddingCat, apply_rot_embed_cat, to_2tuple, use_fused_attn
  from ._builder import build_model_with_cfg
  from ._manipulate import checkpoint_seq
  from ._registry import generate_default_cfgs, register_model
  from ._features_fx import register_notrace_function
label_desc: Use image node_id 1 to create a new VisionTransformerSAM object from inherited base classes: nn with example: obj = VisionTransformerSAM(img_size, patch_size, in_chans, num_classes, embed_dim, depth, num_heads, mlp_ratio, qkv_bias, qk_norm, init_values, pre_norm, drop_rate, pos_drop_rate, patch_drop_rate, proj_drop_rate, attn_drop_rate, drop_path_rate, weight_init, embed_layer, norm_layer, act_layer, block_fn, mlp_layer, use_abs_pos, use_rel_pos, use_rope, window_size, global_attn_indexes, neck_chans, global_pool, head_hidden_size, ref_feat_shape)