class_docstr: Shifted Window Multihead Self-Attention Module.
  
  Args:
      embed_dims (int): Number of input channels.
      num_heads (int): Number of attention heads.
      window_size (int): The height and width of the window.
      shift_size (int, optional): The shift step of each window towards
          right-bottom. If zero, act as regular window-msa. Defaults to 0.
      qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.
          Default: True
      qk_scale (float | None, optional): Override default qk scale of
          head_dim ** -0.5 if set. Defaults: None.
      attn_drop_rate (float, optional): Dropout ratio of attention weight.
          Defaults: 0.
      proj_drop_rate (float, optional): Dropout ratio of output.
          Defaults: 0.
      dropout_layer (dict, optional): The dropout_layer used before output.
          Defaults: dict(type='DropPath', drop_prob=0.).
      init_cfg (dict, optional): The extra config for initialization.
          Default: None.
class_docstr_tok: ["Shifted","Window","Multihead","Self-Attention","Module",".","Args",":","embed_dims","(","int",")",":","Number","of","input","channels",".","num_heads","(","int",")",":","Number","of","attention","heads",".","window_size","(","int",")",":","The","height","and","width","of","the","window",".","shift_size","(","int",",","optional",")",":","The","shift","step","of","each","window","towards","right-bottom",".","If","zero",",","act","as","regular","window-msa",".","Defaults","to","0.","qkv_bias","(","bool",",","optional",")",":","If","True",",","add","a","learnable","bias","to","q",",","k",",","v.","Default",":","True","qk_scale","(","float","|","None",",","optional",")",":","Override","default","qk","scale","of","head_dim","*","*","-0.5","if","set",".","Defaults",":","None",".","attn_drop_rate","(","float",",","optional",")",":","Dropout","ratio","of","attention","weight",".","Defaults",":","0.","proj_drop_rate","(","float",",","optional",")",":","Dropout","ratio","of","output",".","Defaults",":","0.","dropout_layer","(","dict",",","optional",")",":","The","dropout_layer","used","before","output",".","Defaults",":","dict","(","type='DropPath","'",",","drop_prob=0.",")",".","init_cfg","(","dict",",","optional",")",":","The","extra","config","for","initialization",".","Default",":","None","."]
code: def forward(self, query, hw_shape):
      B, L, C = query.shape
      H, W = hw_shape
      assert L == H * W, "input feature has wrong size"
      query = query.view(B, H, W, C)
  
      # pad feature maps to multiples of window size
      pad_r = (
          self.window_size - W % self.window_size
      ) % self.window_size
      pad_b = (
          self.window_size - H % self.window_size
      ) % self.window_size
      query = F.pad(query, (0, 0, 0, pad_r, 0, pad_b))
      H_pad, W_pad = query.shape[1], query.shape[2]
  
      # cyclic shift
      if self.shift_size > 0:
          shifted_query = torch.roll(
              query,
              shifts=(-self.shift_size, -self.shift_size),
              dims=(1, 2),
          )
  
          # calculate attention mask for SW-MSA
          img_mask = torch.zeros(
              (1, H_pad, W_pad, 1), device=query.device
          )
          h_slices = (
              slice(0, -self.window_size),
              slice(-self.window_size, -self.shift_size),
              slice(-self.shift_size, None),
          )
          w_slices = (
              slice(0, -self.window_size),
              slice(-self.window_size, -self.shift_size),
              slice(-self.shift_size, None),
          )
          cnt = 0
          for h in h_slices:
              for w in w_slices:
                  img_mask[:, h, w, :] = cnt
                  cnt += 1
  
          # nW, window_size, window_size, 1
          mask_windows = self.window_partition(img_mask)
          mask_windows = mask_windows.view(
              -1, self.window_size * self.window_size
          )
          attn_mask = mask_windows.unsqueeze(
              1
          ) - mask_windows.unsqueeze(2)
          attn_mask = attn_mask.masked_fill(
              attn_mask != 0, float(-100.0)
          ).masked_fill(attn_mask == 0, float(0.0))
      else:
          shifted_query = query
          attn_mask = None
  
      # nW*B, window_size, window_size, C
      query_windows = self.window_partition(shifted_query)
      # nW*B, window_size*window_size, C
      query_windows = query_windows.view(-1, self.window_size**2, C)
  
      # W-MSA/SW-MSA (nW*B, window_size*window_size, C)
      attn_windows = self.w_msa(query_windows, mask=attn_mask)
  
      # merge windows
      attn_windows = attn_windows.view(
          -1, self.window_size, self.window_size, C
      )
  
      # B H' W' C
      shifted_x = self.window_reverse(attn_windows, H_pad, W_pad)
      # reverse cyclic shift
      if self.shift_size > 0:
          x = torch.roll(
              shifted_x,
              shifts=(self.shift_size, self.shift_size),
              dims=(1, 2),
          )
      else:
          x = shifted_x
  
      if pad_r > 0 or pad_b:
          x = x[:, :H, :W, :].contiguous()
  
      x = x.view(B, H * W, C)
  
      x = self.drop(x)
      return x
  
imports: import warnings
  from collections import OrderedDict
  from copy import deepcopy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from mmcv.cnn import build_norm_layer
  from mmcv.cnn.bricks.transformer import FFN, build_dropout
  from mmengine.logging import MMLogger
  from mmengine.model import BaseModule, ModuleList
  from mmengine.model.weight_init import constant_init, trunc_normal_, trunc_normal_init
  from mmengine.runner.checkpoint import CheckpointLoader
  from mmengine.utils import to_2tuple
  from mmdet.registry import MODELS
  from ..layers import PatchEmbed, PatchMerging
label_desc: Use image node_id 2 for calling the ShiftWindowMSA obj's underlying member method code with example usage: obj.forward(query, hw_shape) and returns: x