class_docstr: Multi-head attention layer.
  
  Args:
      dim (`int`):
          number of features
      num_heads (`int=8`):
          number of heads
      qkv_bias (`bool=False`):
          if add bias to qkv layer
      qk_scale (`float=None`):
          number to scale qk
      attn_drop_ratio (`float=0.`):
          drop rate of attention layer
      proj_drop_ratio (`float=0.`):
          drop rate of projection layer
      with_qkv (`bool=True`):
          if use qkv layer
  
  Example:
      >>> import torch
      >>> from towhee.models.layers.attention import MultiHeadAttention
      >>>
      >>> test_shape = (1, 196+1, 768)  # shape of output from patch_embed
      >>> input_x = torch.rand(test_shape)
      >>> model = MultiHeadAttention(dim=test_shape[2])
      >>> out = model.forward(input_x)
      >>> print(out.shape)
      torch.Size([1, 197, 768])
code: def forward(self, x):
      batch_size, new_num_patch, dim = x.shape
  
      qkv = (
          self.qkv(x)
          .reshape(
              batch_size,
              new_num_patch,
              3,
              self.num_heads,
              self.head_dim,
          )
          .permute(2, 0, 3, 1, 4)
      )
      q, k, v = qkv[0], qkv[1], qkv[2]
  
      self.save_v(v)
  
      self.matmul1([q, k])
      attn = (q @ k.transpose(-2, -1)) * self.scale
      attn = attn.softmax(dim=-1)
      attn = self.attn_drop(attn)
      self.save_attn(attn)
      attn.register_hook(self.save_attn_gradients)
  
      self.matmul2([attn, v])
      x = (
          (attn @ v)
          .transpose(1, 2)
          .reshape(batch_size, new_num_patch, dim)
      )
      x = self.proj(x)
      x = self.proj_drop(x)
      return x
  
imports: import os
  import math
  import numpy
  import torch
  from torch import nn
  from towhee.models.layers.layers_with_relprop import Einsum, Linear, Dropout, Softmax
label_desc: Use image node_id 12 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(x) and returns: x