class_docstr: A Swin Transformer V2 Stage.
      
docstr: Args:
      dim: Number of input channels.
      out_dim: Number of output channels.
      input_resolution: Input resolution.
      depth: Number of blocks.
      num_heads: Number of attention heads.
      window_size: Local window size.
      downsample: Use downsample layer at start of the block.
      mlp_ratio: Ratio of mlp hidden dim to embedding dim.
      qkv_bias: If True, add a learnable bias to query, key, value.
      proj_drop: Projection dropout rate
      attn_drop: Attention dropout rate.
      drop_path: Stochastic depth rate.
      norm_layer: Normalization layer.
      pretrained_window_size: Local window size in pretraining.
      output_nchw: Output tensors on NCHW format instead of NHWC.
docstr_tok: ["Args",":","dim",":","Number","of","input","channels",".","out_dim",":","Number","of","output","channels",".","input_resolution",":","Input","resolution",".","depth",":","Number","of","blocks",".","num_heads",":","Number","of","attention","heads",".","window_size",":","Local","window","size",".","downsample",":","Use","downsample","layer","at","start","of","the","block",".","mlp_ratio",":","Ratio","of","mlp","hidden","dim","to","embedding","dim",".","qkv_bias",":","If","True",",","add","a","learnable","bias","to","query",",","key",",","value",".","proj_drop",":","Projection","dropout","rate","attn_drop",":","Attention","dropout","rate",".","drop_path",":","Stochastic","depth","rate",".","norm_layer",":","Normalization","layer",".","pretrained_window_size",":","Local","window","size","in","pretraining",".","output_nchw",":","Output","tensors","on","NCHW","format","instead","of","NHWC","."]
code: def __init__(
      self,
      dim: int,
      out_dim: int,
      input_resolution: _int_or_tuple_2_t,
      depth: int,
      num_heads: int,
      window_size: _int_or_tuple_2_t,
      downsample: bool = False,
      mlp_ratio: float = 4.0,
      qkv_bias: bool = True,
      proj_drop: float = 0.0,
      attn_drop: float = 0.0,
      drop_path: float = 0.0,
      norm_layer: nn.Module = nn.LayerNorm,
      pretrained_window_size: _int_or_tuple_2_t = 0,
      output_nchw: bool = False,
  ) -> None:
      """
      Args:
          dim: Number of input channels.
          out_dim: Number of output channels.
          input_resolution: Input resolution.
          depth: Number of blocks.
          num_heads: Number of attention heads.
          window_size: Local window size.
          downsample: Use downsample layer at start of the block.
          mlp_ratio: Ratio of mlp hidden dim to embedding dim.
          qkv_bias: If True, add a learnable bias to query, key, value.
          proj_drop: Projection dropout rate
          attn_drop: Attention dropout rate.
          drop_path: Stochastic depth rate.
          norm_layer: Normalization layer.
          pretrained_window_size: Local window size in pretraining.
          output_nchw: Output tensors on NCHW format instead of NHWC.
      """
      super().__init__()
      self.dim = dim
      self.input_resolution = input_resolution
      self.output_resolution = (
          tuple(i // 2 for i in input_resolution)
          if downsample
          else input_resolution
      )
      self.depth = depth
      self.output_nchw = output_nchw
      self.grad_checkpointing = False
      window_size = to_2tuple(window_size)
      shift_size = tuple([w // 2 for w in window_size])
  
      # patch merging / downsample layer
      if downsample:
          self.downsample = PatchMerging(
              dim=dim, out_dim=out_dim, norm_layer=norm_layer
          )
      else:
          assert dim == out_dim
          self.downsample = nn.Identity()
  
      # build blocks
      self.blocks = nn.ModuleList(
          [
              SwinTransformerV2Block(
                  dim=out_dim,
                  input_resolution=self.output_resolution,
                  num_heads=num_heads,
                  window_size=window_size,
                  shift_size=0 if (i % 2 == 0) else shift_size,
                  mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias,
                  proj_drop=proj_drop,
                  attn_drop=attn_drop,
                  drop_path=drop_path[i]
                  if isinstance(drop_path, list)
                  else drop_path,
                  norm_layer=norm_layer,
                  pretrained_window_size=pretrained_window_size,
              )
              for i in range(depth)
          ]
      )
  
imports: import math
  from typing import Callable, Optional, Tuple, Union, Set, Dict
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import PatchEmbed, Mlp, DropPath, to_2tuple, trunc_normal_, _assert, ClassifierHead, resample_patch_embed
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_function
  from ._registry import generate_default_cfgs, register_model, register_model_deprecations
label_desc: Use image node_id 1 to create a new SwinTransformerV2Stage object from inherited base classes: nn with example: obj = SwinTransformerV2Stage(dim, out_dim, input_resolution, depth, num_heads, window_size, downsample, mlp_ratio, qkv_bias, proj_drop, attn_drop, drop_path, norm_layer, pretrained_window_size, output_nchw)