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,
      logit_scale: Optional[torch.Tensor] = None,
      training: bool = True,
  ) -> Tensor:
      """
      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.
          logit_scale (Tensor[out_dim], optional): Logit scale of cosine attention for Swin Transformer V2. Default: None.
          training (bool, optional): Training flag used by the dropout parameters. Default: True.
      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
      if logit_scale is not None and qkv_bias is not None:
          qkv_bias = qkv_bias.clone()
          length = qkv_bias.numel() // 3
          qkv_bias[length : 2 * length].zero_()
      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]
      if logit_scale is not None:
          # cosine attention
          attn = F.normalize(q, dim=-1) @ F.normalize(
              k, dim=-1
          ).transpose(-2, -1)
          logit_scale = torch.clamp(
              logit_scale, max=math.log(100.0)
          ).exp()
          attn = attn * logit_scale
      else:
          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, training=training)
  
      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, training=training)
  
      # 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: import math
  from functools import partial
  from typing import Any, Callable, List, Optional
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
  from ..ops.misc import MLP, Permute
  from ..ops.stochastic_depth import StochasticDepth
  from ..transforms._presets import ImageClassification, InterpolationMode
  from ..utils import _log_api_usage_once
  from ._api import register_model, Weights, WeightsEnum
  from ._meta import _IMAGENET_CATEGORIES
  from ._utils import _ovewrite_named_param, handle_legacy_interface
label_desc: Use image node_id 3 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, logit_scale, training) and returns: x