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
class_docstr_tok: ["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"]
code: def __init__(
      self,
      dim,
      window_size,
      num_heads,
      qkv_bias=True,
      qk_scale=None,
      attn_drop=0.0,
      proj_drop=0.0,
  ):
      super(WindowAttention, self).__init__()
      self.dim = dim
      self.window_size = window_size  # Wh, Ww
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = qk_scale or head_dim**-0.5
  
      # define a parameter table of relative position bias
      self.relative_position_bias_table = nn.Parameter(
          torch.zeros(
              (2 * window_size[0] - 1) * (2 * window_size[1] - 1),
              num_heads,
          )
      )  # 2*Wh-1 * 2*Ww-1, nH
  
      # get pair-wise relative position index for each token inside the window
      coords_h = torch.arange(self.window_size[0])
      coords_w = torch.arange(self.window_size[1])
      coords = torch.stack(
          torch.meshgrid([coords_h, coords_w])
      )  # 2, Wh, Ww
      coords_flatten = torch.flatten(coords, 1)  # 2 Wh*Ww
      relative_coords = (
          coords_flatten[:, :, None] - coords_flatten[:, None, :]
      )  # 2, Wh*Ww, Wh*Ww
      relative_coords = relative_coords.permute(
          1, 2, 0
      ).contiguous()  # Wh*Ww, Wh*Ww, 2
      relative_coords[:, :, 0] += (
          self.window_size[0] - 1
      )  # shift to start from 0
      relative_coords[:, :, 1] += self.window_size[1] - 1
      relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
      relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
      self.register_buffer(
          "relative_position_index", relative_position_index
      )
  
      self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj = nn.Linear(dim, dim)
      self.proj_drop = nn.Dropout(proj_drop)
  
      trunc_normal_(self.relative_position_bias_table, std=0.02)
      self.softmax = nn.Softmax(dim=-1)
  
imports: import os
  import logging
  import torch
  import torch.nn
  import torch.nn.functional
  from functools import partial
  import torch.distributed
  from torch.nn.init import trunc_normal_
  from megatron.model.transformer import DropPath
  from megatron import get_args
  from megatron.model import LayerNorm
  import numpy
  from math import sqrt
label_desc: Use image node_id 1 to create a new WindowAttention object from inherited base classes: nn with example: obj = WindowAttention(dim, window_size, num_heads, qkv_bias, qk_scale, attn_drop, proj_drop)