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 flops(self, N):
      # calculate flops for 1 window with token length of N
      flops = 0
      # qkv = self.qkv(x)
      flops += N * self.dim * 3 * self.dim
      # attn = (q @ k.transpose(-2, -1))
      flops += self.num_heads * N * (self.dim // self.num_heads) * N
      #  x = (attn @ v)
      flops += self.num_heads * N * N * (self.dim // self.num_heads)
      # x = self.proj(x)
      flops += N * self.dim * self.dim
      return flops
  
imports: import torch
  import torch.nn
  import torch.utils.checkpoint
  from timm.models.layers import DropPath, to_2tuple, trunc_normal_
  from math import sqrt
  from megatron import get_args
  from functools import partial
label_desc: Use image node_id 4 for calling the WindowAttention obj's underlying member method code with example usage: obj.flops(N) and returns: flops