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
      attn_drop_ratio (float, optional): Dropout ratio of attention weight. Default: 0.0
      proj_drop_ratio (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","attn_drop_ratio","(","float",",","optional",")",":","Dropout","ratio","of","attention","weight",".","Default",":","0.0","proj_drop_ratio","(","float",",","optional",")",":","Dropout","ratio","of","output",".","Default",":","0.0"]
docstr: Args:
      x: input features with shape of (num_windows*B, Mh*Mw, C)
      mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
      training: whether training mode
code: def call(self, x, mask=None, training=None):
      """
      Args:
          x: input features with shape of (num_windows*B, Mh*Mw, C)
          mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
          training: whether training mode
      """
      # [batch_size*num_windows, Mh*Mw, total_embed_dim]
      B_, N, C = x.shape
  
      # qkv(): -> [batch_size*num_windows, Mh*Mw, 3 * total_embed_dim]
      qkv = self.qkv(x)
      # reshape: -> [batch_size*num_windows, Mh*Mw, 3, num_heads, embed_dim_per_head]
      qkv = tf.reshape(
          qkv, [B_, N, 3, self.num_heads, C // self.num_heads]
      )
      # transpose: -> [3, batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
      qkv = tf.transpose(qkv, [2, 0, 3, 1, 4])
      # [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
      q, k, v = qkv[0], qkv[1], qkv[2]
  
      # transpose: -> [batch_size*num_windows, num_heads, embed_dim_per_head, Mh*Mw]
      # multiply -> [batch_size*num_windows, num_heads, Mh*Mw, Mh*Mw]
      attn = tf.matmul(a=q, b=k, transpose_b=True) * self.scale
  
      # relative_position_bias(reshape): [Mh*Mw*Mh*Mw,nH] -> [Mh*Mw,Mh*Mw,nH]
      relative_position_bias = tf.gather(
          self.relative_position_bias_table,
          tf.reshape(self.relative_position_index, [-1]),
      )
      relative_position_bias = tf.reshape(
          relative_position_bias,
          [
              self.window_size[0] * self.window_size[1],
              self.window_size[0] * self.window_size[1],
              -1,
          ],
      )
      relative_position_bias = tf.transpose(
          relative_position_bias, [2, 0, 1]
      )  # [nH, Mh*Mw, Mh*Mw]
      attn = attn + tf.expand_dims(relative_position_bias, 0)
  
      if mask is not None:
          # mask: [nW, Mh*Mw, Mh*Mw]
          nW = mask.shape[0]  # num_windows
          # attn(reshape): [batch_size, num_windows, num_heads, Mh*Mw, Mh*Mw]
          # mask(expand_dim): [1, nW, 1, Mh*Mw, Mh*Mw]
          attn = tf.reshape(
              attn, [B_ // nW, nW, self.num_heads, N, N]
          ) + tf.expand_dims(tf.expand_dims(mask, 1), 0)
          attn = tf.reshape(attn, [-1, self.num_heads, N, N])
  
      attn = tf.nn.softmax(attn, axis=-1)
      attn = self.attn_drop(attn, training=training)
  
      # multiply -> [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
      x = tf.matmul(attn, v)
      # transpose: -> [batch_size*num_windows, Mh*Mw, num_heads, embed_dim_per_head]
      x = tf.transpose(x, [0, 2, 1, 3])
      # reshape: -> [batch_size*num_windows, Mh*Mw, total_embed_dim]
      x = tf.reshape(x, [B_, N, C])
  
      x = self.proj(x)
      x = self.proj_drop(x, training=training)
      return x
  
imports: import tensorflow
  from tensorflow.keras import Model, layers, initializers
  import numpy
label_desc: Use image node_id 3 for calling the WindowAttention obj's underlying member method code with example usage: obj.call(x, mask, training) and returns: x