class_docstr: Window based multi-head self attention (W-MSA) module with relative position bias.
  It supports both of shifted and non-shifted window.
  
  Args:
      dim (int): Number of input channels.
      window_size (tuple[int]): The height and width of the window.
      num_heads (int): Number of attention heads.
      qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
      qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
      attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
      proj_drop (float, optional): Dropout ratio of output. Default: 0.0
docstr: Args:
      x: input features with shape of (num_windows*B, N, C)
      mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
docstr_tok: ["Args",":","x",":","input","features","with","shape","of","(","num_windows","*","B",",","N",",","C",")","mask",":","(","0\/-inf",")","mask","with","shape","of","(","num_windows",",","Wh","*","Ww",",","Wh","*","Ww",")","or","None"]
code: def forward(self, x, mask=None):
      """
      Args:
          x: input features with shape of (num_windows*B, N, C)
          mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
      """
      B_, N, C = x.shape
      qkv = (
          self.qkv(x)
          .reshape(B_, N, 3, self.num_heads, C // self.num_heads)
          .permute(2, 0, 3, 1, 4)
      )
      q, k, v = (
          qkv[0],
          qkv[1],
          qkv[2],
      )  # make torchscript happy (cannot use tensor as tuple)
  
      q = q * self.scale
      attn = q @ k.transpose(-2, -1)
  
      relative_position_bias = self.relative_position_bias_table[
          self.relative_position_index[:N, :N].reshape(-1)
      ].reshape(
          N, N, -1
      )  # N, N,nH
      relative_position_bias = relative_position_bias.permute(
          2, 0, 1
      ).contiguous()  # nH, N, N
      attn = attn + relative_position_bias.unsqueeze(0)
  
      if mask is not None:
          nW = mask.shape[0]
          attn = attn.view(
              B_ // nW, nW, self.num_heads, N, N
          ) + mask.unsqueeze(1).unsqueeze(0)
          attn = attn.view(-1, self.num_heads, N, N)
          attn = self.softmax(attn)
      else:
          attn = self.softmax(attn)
  
      attn = self.attn_drop(attn)
  
      x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
      x = self.proj(x)
      x = self.proj_drop(x)
      return x
  
imports: import math
  from functools import lru_cache
  from typing import List, Optional
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from vissl.config import AttrDict
  from vissl.models.model_helpers import DropPath, rearrange, to_2tuple, trunc_normal_
  from vissl.models.trunks import register_model_trunk
label_desc: Use image node_id 2 for calling the WindowAttention obj's underlying member method code with example usage: obj.forward(x, mask) and returns: x