class_docstr: Attention module
  Attention module for ViT, here q, k, v are assumed the same.
  The qkv mappings are stored as one single param.
  Attributes:
      num_heads: number of heads
      attn_head_size: feature dim of single head
      all_head_size: feature dim of all heads
      qkv: a nn.Linear for q, k, v mapping
      scales: 1 / sqrt(single_head_feature_dim)
      out: projection of multi-head attention
      attn_dropout: dropout for attention
      proj_dropout: final dropout before output
      softmax: softmax op for attention
code: def __init__(
      self,
      embed_dim,
      num_heads,
      attn_head_size=None,
      qkv_bias=True,
      dropout=0.0,
      attention_dropout=0.0,
  ):
      super().__init__()
  
      assert isinstance(
          embed_dim, int
      ), f"Expected the type of `embed_dim` to be {int}, but received {type(embed_dim)}."
      assert isinstance(
          num_heads, int
      ), f"Expected the type of `num_heads` to be {int}, but received {type(num_heads)}."
  
      assert (
          embed_dim > 0
      ), f"Expected `embed_dim` to be greater than 0, but received {embed_dim}"
      assert (
          num_heads > 0
      ), f"Expected `num_heads` to be greater than 0, but received {num_heads}"
  
      self.embed_dim = embed_dim
      self.num_heads = num_heads
  
      if attn_head_size is not None:
          assert isinstance(attn_head_size, int), (
              f"Expected the type of `attn_head_size` to be {int}, "
              f"but received {type(attn_head_size)}."
          )
          assert attn_head_size > 0, (
              f"Expected `attn_head_size` to be greater than 0,"
              f" but received {attn_head_size}."
          )
          self.attn_head_size = attn_head_size
      else:
          self.attn_head_size = embed_dim // num_heads
          assert self.attn_head_size * num_heads == embed_dim, (
              f"`embed_dim` must be divisible by `num_heads`,"
              f" but received embed_dim={embed_dim}, num_heads={num_heads}."
          )
  
      self.all_head_size = self.attn_head_size * num_heads
  
      w_attr_1, b_attr_1 = self._init_weights()
      self.qkv = nn.Linear(
          embed_dim,
          self.all_head_size * 3,  # weights for q, k, and v
          weight_attr=w_attr_1,
          bias_attr=b_attr_1 if qkv_bias else False,
      )
  
      self.scales = self.attn_head_size**-0.5
  
      w_attr_2, b_attr_2 = self._init_weights()
      self.out = nn.Linear(
          self.all_head_size,
          embed_dim,
          weight_attr=w_attr_2,
          bias_attr=b_attr_2,
      )
  
      self.attn_dropout = nn.Dropout(attention_dropout)
      self.proj_dropout = nn.Dropout(dropout)
      self.softmax = nn.Softmax(axis=-1)
  
imports: import copy
  import paddle
  import paddle.nn
  from disco_diffusion_ernievil_base.vit_b_16x.ernievil2.transformers.droppath import DropPath
label_desc: Use image node_id 1 to create a new Attention object from inherited base classes: nn with example: obj = Attention(embed_dim, num_heads, attn_head_size, qkv_bias, dropout, attention_dropout)