class_docstr: Self-attention module in Swin Transformer
      
code: def __init__(
      self, input_dim, output_dim, head_dim, window_size, type
  ):
      super(WMSA, self).__init__()
      self.input_dim = input_dim
      self.output_dim = output_dim
      self.head_dim = head_dim
      self.scale = self.head_dim**-0.5
      self.n_heads = input_dim // head_dim
      self.window_size = window_size
      self.type = type
      self.embedding_layer = nn.Linear(
          self.input_dim, 3 * self.input_dim, bias=True
      )
  
      self.relative_position_params = nn.Parameter(
          torch.zeros(
              (2 * window_size - 1) * (2 * window_size - 1),
              self.n_heads,
          )
      )
  
      self.linear = nn.Linear(self.input_dim, self.output_dim)
  
      trunc_normal_(self.relative_position_params, std=0.02)
      self.relative_position_params = torch.nn.Parameter(
          self.relative_position_params.view(
              2 * window_size - 1, 2 * window_size - 1, self.n_heads
          )
          .transpose(1, 2)
          .transpose(0, 1)
      )
  
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 1 to create a new WMSA object from inherited base classes: nn with example: obj = WMSA(input_dim, output_dim, head_dim, window_size, type)