class_docstr: dim:
  num_heads:
  qkv_bias:
  qk_scale:
  attn_drop:
  proj_drop:
class_docstr_tok: ["dim",":","num_heads",":","qkv_bias",":","qk_scale",":","attn_drop",":","proj_drop",":"]
code: def __init__(
      self,
      dim,
      num_heads=8,
      qkv_bias=False,
      qk_scale=None,
      attn_drop=0.0,
      proj_drop=0.0,
      initialize="random",
  ):
      super().__init__()
      self.num_heads = num_heads
      head_dim = dim // num_heads
      # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
      self.scale = qk_scale or head_dim**-0.5
      self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
      self.proj = nn.Linear(dim, dim)
      if initialize == "zeros":
          self.qkv.weight.data.fill_(0)
          self.qkv.bias.data.fill_(0)
          # fill proj weight with 1 here to improve training dynamics. Otherwise temporal attention inputs
          # are multiplied by 0*0, which is hard for the model to move out of.
          self.proj.weight.data.fill_(1)
          self.proj.bias.data.fill_(0)
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj_drop = nn.Dropout(proj_drop)
  
imports: from collections import OrderedDict
  from functools import partial
  import torch
  from einops import rearrange, repeat
  from timm.models.layers import DropPath, trunc_normal_
  from torch import einsum, nn
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from towhee.models.utils.init_vit_weights import init_vit_weights
  import logging
label_desc: Use image node_id 1 to create a new VarAttention object from inherited base classes: nn with example: obj = VarAttention(dim, num_heads, qkv_bias, qk_scale, attn_drop, proj_drop, initialize)