class_docstr: Self-attention module in Swin Transformer
      
docstr: Forward pass of Window Multi-head Self-attention module.
  Args:
      x: input tensor with shape of [b h w c];
      attn_mask: attention mask, fill -inf where the value is True;
  Returns:
      output: tensor shape [b h w c]
docstr_tok: ["Forward","pass","of","Window","Multi-head","Self-attention","module",".","Args",":","x",":","input","tensor","with","shape","of","[","b","h","w","c","]",";","attn_mask",":","attention","mask",",","fill","-inf","where","the","value","is","True",";","Returns",":","output",":","tensor","shape","[","b","h","w","c","]"]
code: def forward(self, x):
      """Forward pass of Window Multi-head Self-attention module.
      Args:
          x: input tensor with shape of [b h w c];
          attn_mask: attention mask, fill -inf where the value is True;
      Returns:
          output: tensor shape [b h w c]
      """
      if self.type != "W":
          x = torch.roll(
              x,
              shifts=(
                  -(self.window_size // 2),
                  -(self.window_size // 2),
              ),
              dims=(1, 2),
          )
  
      x = rearrange(
          x,
          "b (w1 p1) (w2 p2) c -> b w1 w2 p1 p2 c",
          p1=self.window_size,
          p2=self.window_size,
      )
      h_windows = x.size(1)
      w_windows = x.size(2)
      # square validation
      # assert h_windows == w_windows
  
      x = rearrange(
          x,
          "b w1 w2 p1 p2 c -> b (w1 w2) (p1 p2) c",
          p1=self.window_size,
          p2=self.window_size,
      )
      qkv = self.embedding_layer(x)
      q, k, v = rearrange(
          qkv, "b nw np (threeh c) -> threeh b nw np c", c=self.head_dim
      ).chunk(3, dim=0)
      sim = torch.einsum("hbwpc,hbwqc->hbwpq", q, k) * self.scale
      # Adding learnable relative embedding
      sim = sim + rearrange(
          self.relative_embedding(), "h p q -> h 1 1 p q"
      )
      # Using Attn Mask to distinguish different subwindows.
      if self.type != "W":
          attn_mask = self.generate_mask(
              h_windows,
              w_windows,
              self.window_size,
              shift=self.window_size // 2,
          )
          sim = sim.masked_fill_(attn_mask, float("-inf"))
  
      probs = nn.functional.softmax(sim, dim=-1)
      output = torch.einsum("hbwij,hbwjc->hbwic", probs, v)
      output = rearrange(output, "h b w p c -> b w p (h c)")
      output = self.linear(output)
      output = rearrange(
          output,
          "b (w1 w2) (p1 p2) c -> b (w1 p1) (w2 p2) c",
          w1=h_windows,
          p1=self.window_size,
      )
  
      if self.type != "W":
          output = torch.roll(
              output,
              shifts=(self.window_size // 2, self.window_size // 2),
              dims=(1, 2),
          )
  
      return output
  
imports: import numpy
  import torch
  import torch.nn
  from einops import rearrange
  from einops.layers.torch import Rearrange
  from timm.models.layers import trunc_normal_, DropPath
label_desc: Use image node_id 3 for calling the WMSA obj's underlying member method code with example usage: obj.forward(x) and returns: output