class_docstr: This layer applies a multi-head self- or cross-attention as described in
  `Attention is all you need <https://arxiv.org/abs/1706.03762>`_ paper
  
  Args:
      embed_dim (int): :math:`C_{in}` from an expected input of size :math:`(N, S, C_{in})`
      num_heads (int): Number of heads in multi-head attention
      attn_dropout (Optional[float]): Attention dropout. Default: 0.0
      bias (Optional[bool]): Use bias or not. Default: ``True``
  
  Shape:
      - Input:
         - Query tensor (x_q) :math:`(N, S, C_{in})` where :math:`N` is batch size, :math:`S` is number of source tokens,
      and :math:`C_{in}` is input embedding dim
         - Optional Key-Value tensor (x_kv) :math:`(N, T, C_{in})` where :math:`T` is number of target tokens
      - Output: same shape as the input
code: def __init__(
      self,
      embed_dim: int,
      num_heads: int,
      attn_dropout: Optional[float] = 0.0,
      bias: Optional[bool] = True,
      output_dim: Optional[int] = None,
      coreml_compatible: Optional[bool] = False,
      *args,
      **kwargs,
  ) -> None:
      if output_dim is None:
          output_dim = embed_dim
      super().__init__()
      if embed_dim % num_heads != 0:
          logger.error(
              "Embedding dim must be divisible by number of heads in {}. Got: embed_dim={} and num_heads={}".format(
                  self.__class__.__name__, embed_dim, num_heads
              )
          )
  
      self.qkv_proj = LinearLayer(
          in_features=embed_dim, out_features=3 * embed_dim, bias=bias
      )
  
      self.attn_dropout = Dropout(p=attn_dropout)
      self.out_proj = LinearLayer(
          in_features=embed_dim, out_features=output_dim, bias=bias
      )
  
      self.head_dim = embed_dim // num_heads
      self.scaling = self.head_dim**-0.5
      self.softmax = nn.Softmax(dim=-1)
      self.num_heads = num_heads
      self.embed_dim = embed_dim
      self.coreml_compatible = coreml_compatible
      self.use_separate_proj_weight = embed_dim != output_dim
  
file_path: ml-cvnets/cvnets/layers/multi_head_attention.py
imports: from typing import Optional
  import torch
  from torch import Tensor, nn
  from torch.nn import functional
  from cvnets.layers.base_layer import BaseLayer
  from cvnets.layers.dropout import Dropout
  from cvnets.layers.linear_layer import LinearLayer
  from utils import logger
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: BaseLayer with example: obj = MultiHeadAttention(embed_dim, num_heads, attn_dropout, bias, output_dim, coreml_compatible)