code: def __init__(
      self,
      dim,
      num_heads=8,
      qkv_bias=True,
      attn_drop=0.0,
      proj_drop=0.0,
      window_size=None,
      attn_head_dim=None,
  ):
      super().__init__()
      self.num_heads = num_heads
      head_dim = dim // num_heads
      if attn_head_dim is not None:
          head_dim = attn_head_dim
      all_head_dim = head_dim * self.num_heads
      self.scale = head_dim**-0.5
  
      self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
      if qkv_bias:
          self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
          self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
      else:
          self.q_bias = None
          self.v_bias = None
  
      if window_size:
          self.window_size = window_size
          self.num_relative_distance = (2 * window_size[0] - 1) * (
              2 * window_size[1] - 1
          ) + 3
          self.relative_position_bias_table = nn.Parameter(
              torch.zeros(self.num_relative_distance, num_heads)
          )  # 2*Wh-1 * 2*Ww-1, nH
          # cls to token & token 2 cls & cls to cls
  
          # get pair-wise relative position index for each token inside the window
          coords_h = torch.arange(window_size[0])
          coords_w = torch.arange(window_size[1])
          coords = torch.stack(
              torch.meshgrid([coords_h, coords_w])
          )  # 2, Wh, Ww
          coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
          relative_coords = (
              coords_flatten[:, :, None] - coords_flatten[:, None, :]
          )  # 2, Wh*Ww, Wh*Ww
          relative_coords = relative_coords.permute(
              1, 2, 0
          ).contiguous()  # Wh*Ww, Wh*Ww, 2
          relative_coords[:, :, 0] += (
              window_size[0] - 1
          )  # shift to start from 0
          relative_coords[:, :, 1] += window_size[1] - 1
          relative_coords[:, :, 0] *= 2 * window_size[1] - 1
          relative_position_index = torch.zeros(
              size=(window_size[0] * window_size[1] + 1,) * 2,
              dtype=relative_coords.dtype,
          )
          relative_position_index[1:, 1:] = relative_coords.sum(
              -1
          )  # Wh*Ww, Wh*Ww
          relative_position_index[0, 0:] = (
              self.num_relative_distance - 3
          )
          relative_position_index[0:, 0] = (
              self.num_relative_distance - 2
          )
          relative_position_index[0, 0] = self.num_relative_distance - 1
  
          self.register_buffer(
              "relative_position_index", relative_position_index
          )
      else:
          self.window_size = None
          self.relative_position_bias_table = None
          self.relative_position_index = None
  
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj = nn.Linear(all_head_dim, dim)
      self.proj_drop = nn.Dropout(proj_drop)
  
imports: import logging
  import math
  import numpy
  import random
  from dataclasses import dataclass, field
  from typing import Optional
  from omegaconf import II
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.distributed
  from fairseq.modules import EMAModule, EMAModuleConfig
  from fairseq.dataclass import FairseqDataclass
  from fairseq.models import BaseFairseqModel, register_model
label_desc: Use image node_id 1 to create a new Attention object from inherited base classes: nn with example: obj = Attention(dim, num_heads, qkv_bias, attn_drop, proj_drop, window_size, attn_head_dim)