code: def __init__(
      self,
      dim,
      num_heads=8,
      qkv_bias=False,
      qk_scale=None,
      attn_drop=0.0,
      proj_drop=0.0,
      sr_ratio=1,
  ):
      super().__init__()
      assert (
          dim % num_heads == 0
      ), f"dim {dim} should be divided by num_heads {num_heads}."
  
      self.dim = dim
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = qk_scale or head_dim**-0.5
  
      self.q = nn.Linear(dim, dim, bias=qkv_bias)
      self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj = nn.Linear(dim, dim)
      self.proj_drop = nn.Dropout(proj_drop)
  
      self.sr_ratio = sr_ratio
      if sr_ratio > 1:
          self.sr = nn.Conv2d(
              dim, dim, kernel_size=sr_ratio, stride=sr_ratio
          )
          self.norm = LayerNorm(dim)
  
      self.apply(self._init_weights)
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
  from functools import partial
  from torch.nn.init import trunc_normal_
  from megatron.model.transformer import DropPath
  from megatron.model import LayerNorm
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, qk_scale, attn_drop, proj_drop, sr_ratio)