class_docstr: Window based multi-head self attention (W-MSA) module with relative position bias.
  It supports both of shifted and non-shifted window.
  ref1: https://arxiv.org/pdf/2010.11929.pdf
  ref2: https://arxiv.org/pdf/2106.13230.pdf
  Args:
      dim (`int`):
          Number of input channels, see ref2.
      window_size (`tuple[int]`):
          The temporal length, height and width of the window, see ref2.
      num_heads (`int`):
          Number of attention heads of MSA, see ref1.
      qkv_bias (`bool`):
          If True, add a learnable bias to query, key, value. Default: True
      qk_scale (`float`):
          Override default qk scale of head_dim ** -0.5 if set
      attn_drop (`float`):
          Dropout ratio of attention weight. Default: 0.0
      proj_drop (`float`):
          Dropout ratio of output. Default: 0.0
docstr: Forward function of window attention 3d.
  Args:
      x (`tensor`):
          Input features with shape of (num_windows*b, n, c). See ref2, n = (t * h * w) / (p * m * m).
          num_windows is Wd * Wh * Ww.
          b is batch size. c is dimension or channels size, same as dim param in init function.
      mask (`tensor`):
          Mask with shape of (num_windows, n, n) or None
docstr_tok: ["Forward","function","of","window","attention","3d",".","Args",":","x","(","`","tensor","`",")",":","Input","features","with","shape","of","(","num_windows","*","b",",","n",",","c",")",".","See","ref2",",","n","=","(","t","*","h","*","w",")","\/","(","p","*","m","*","m",")",".","num_windows","is","Wd","*","Wh","*","Ww",".","b","is","batch","size",".","c","is","dimension","or","channels","size",",","same","as","dim","param","in","init","function",".","mask","(","`","tensor","`",")",":","Mask","with","shape","of","(","num_windows",",","n",",","n",")","or","None"]
code: def forward(self, x, mask=None):
      """
      Forward function of window attention 3d.
      Args:
          x (`tensor`):
              Input features with shape of (num_windows*b, n, c). See ref2, n = (t * h * w) / (p * m * m).
              num_windows is Wd * Wh * Ww.
              b is batch size. c is dimension or channels size, same as dim param in init function.
          mask (`tensor`):
              Mask with shape of (num_windows, n, n) 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]  # B_, nH, N, C
  
      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
      )  # Wd*Wh*Ww,Wd*Wh*Ww,nH
      relative_position_bias = relative_position_bias.permute(
          2, 0, 1
      ).contiguous()  # nH, Wd*Wh*Ww, Wd*Wh*Ww
      attn = attn + relative_position_bias.unsqueeze(0)  # B_, nH, N, N
  
      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: from torch import nn
  import torch
  from towhee.models.utils.weight_init import trunc_normal_
label_desc: Use image node_id 2 for calling the WindowAttention3D obj's underlying member method code with example usage: obj.forward(x, mask) and returns: x