code: def shifted_window_attention(
      input: Tensor,
      qkv_weight: Tensor,
      proj_weight: Tensor,
      relative_position_bias: Tensor,
      window_size: List[int],
      num_heads: int,
      shift_size: List[int],
      attention_dropout: float = 0.0,
      dropout: float = 0.0,
      qkv_bias: Optional[Tensor] = None,
      proj_bias: Optional[Tensor] = None,
  ):
      """
      Window based multi-head self attention (W-MSA) module with relative position bias.
      It supports both of shifted and non-shifted window.
      Args:
          input (Tensor[N, H, W, C]): The input tensor or 4-dimensions.
          qkv_weight (Tensor[in_dim, out_dim]): The weight tensor of query, key, value.
          proj_weight (Tensor[out_dim, out_dim]): The weight tensor of projection.
          relative_position_bias (Tensor): The learned relative position bias added to attention.
          window_size (List[int]): Window size.
          num_heads (int): Number of attention heads.
          shift_size (List[int]): Shift size for shifted window attention.
          attention_dropout (float): Dropout ratio of attention weight. Default: 0.0.
          dropout (float): Dropout ratio of output. Default: 0.0.
          qkv_bias (Tensor[out_dim], optional): The bias tensor of query, key, value. Default: None.
          proj_bias (Tensor[out_dim], optional): The bias tensor of projection. Default: None.
      Returns:
          Tensor[N, H, W, C]: The output tensor after shifted window attention.
      """
      B, H, W, C = input.shape
      # pad feature maps to multiples of window size
      pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
      pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
      x = F.pad(input, (0, 0, 0, pad_r, 0, pad_b))
      _, pad_H, pad_W, _ = x.shape
  
      shift_size = shift_size.copy()
      # If window size is larger than feature size, there is no need to shift window
      if window_size[0] >= pad_H:
          shift_size[0] = 0
      if window_size[1] >= pad_W:
          shift_size[1] = 0
  
      # cyclic shift
      if sum(shift_size) > 0:
          x = torch.roll(
              x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2)
          )
  
      # partition windows
      num_windows = (pad_H // window_size[0]) * (
          pad_W // window_size[1]
      )
      x = x.view(
          B,
          pad_H // window_size[0],
          window_size[0],
          pad_W // window_size[1],
          window_size[1],
          C,
      )
      x = x.permute(0, 1, 3, 2, 4, 5).reshape(
          B * num_windows, window_size[0] * window_size[1], C
      )  # B*nW, Ws*Ws, C
  
      # multi-head attention
      qkv = F.linear(x, qkv_weight, qkv_bias)
      qkv = qkv.reshape(
          x.size(0), x.size(1), 3, num_heads, C // num_heads
      ).permute(2, 0, 3, 1, 4)
      q, k, v = qkv[0], qkv[1], qkv[2]
      q = q * (C // num_heads) ** -0.5
      attn = q.matmul(k.transpose(-2, -1))
      # add relative position bias
      attn = attn + relative_position_bias
  
      if sum(shift_size) > 0:
          # generate attention mask
          attn_mask = x.new_zeros((pad_H, pad_W))
          h_slices = (
              (0, -window_size[0]),
              (-window_size[0], -shift_size[0]),
              (-shift_size[0], None),
          )
          w_slices = (
              (0, -window_size[1]),
              (-window_size[1], -shift_size[1]),
              (-shift_size[1], None),
          )
          count = 0
          for h in h_slices:
              for w in w_slices:
                  attn_mask[h[0] : h[1], w[0] : w[1]] = count
                  count += 1
          attn_mask = attn_mask.view(
              pad_H // window_size[0],
              window_size[0],
              pad_W // window_size[1],
              window_size[1],
          )
          attn_mask = attn_mask.permute(0, 2, 1, 3).reshape(
              num_windows, window_size[0] * window_size[1]
          )
          attn_mask = attn_mask.unsqueeze(1) - attn_mask.unsqueeze(2)
          attn_mask = attn_mask.masked_fill(
              attn_mask != 0, float(-100.0)
          ).masked_fill(attn_mask == 0, float(0.0))
          attn = attn.view(
              x.size(0) // num_windows,
              num_windows,
              num_heads,
              x.size(1),
              x.size(1),
          )
          attn = attn + attn_mask.unsqueeze(1).unsqueeze(0)
          attn = attn.view(-1, num_heads, x.size(1), x.size(1))
  
      attn = F.softmax(attn, dim=-1)
      attn = F.dropout(attn, p=attention_dropout)
  
      x = (
          attn.matmul(v)
          .transpose(1, 2)
          .reshape(x.size(0), x.size(1), C)
      )
      x = F.linear(x, proj_weight, proj_bias)
      x = F.dropout(x, p=dropout)
  
      # reverse windows
      x = x.view(
          B,
          pad_H // window_size[0],
          pad_W // window_size[1],
          window_size[0],
          window_size[1],
          C,
      )
      x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, pad_H, pad_W, C)
  
      # reverse cyclic shift
      if sum(shift_size) > 0:
          x = torch.roll(
              x, shifts=(shift_size[0], shift_size[1]), dims=(1, 2)
          )
  
      # unpad features
      x = x[:, :H, :W, :].contiguous()
      return x
  
imports: from typing import List, Optional
  import torch
  from torch import Tensor, nn
  from torch.nn import functional
  from cvnets.layers import Dropout, LinearLayer, StochasticDepth, get_normalization_layer
  from cvnets.layers.activation import build_activation_layer
  from cvnets.modules import BaseModule
label_desc: Use image node_id 2 for calling a global function with example usage: shifted_window_attention(input, qkv_weight, proj_weight, relative_position_bias, window_size, num_heads, shift_size, attention_dropout, dropout, qkv_bias, proj_bias) and returns: x