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 forward_tracing(
      self,
      x_q: Tensor,
      x_kv: Optional[Tensor] = None,
      key_padding_mask: Optional[Tensor] = None,
      attn_mask: Optional[Tensor] = None,
  ) -> Tensor:
      if x_kv is None:
          # [N, S, C] --> # [N, S, 3C] Here, T=S
          qkv = self.qkv_proj(x_q)
          # # [N, S, 3C] --> # [N, S, C] x 3
          query, key, value = torch.chunk(qkv, chunks=3, dim=-1)
      else:
          # [N, S, C]
          query = F.linear(
              x_q,
              weight=self.qkv_proj.weight[: self.embed_dim, ...],
              bias=self.qkv_proj.bias[: self.embed_dim]
              if self.qkv_proj.bias is not None
              else None,
          )
  
          # [N, T, C] --> [N, T, 2C]
          kv = F.linear(
              x_kv,
              weight=self.qkv_proj.weight[self.embed_dim :, ...],
              bias=self.qkv_proj.bias[self.embed_dim :]
              if self.qkv_proj.bias is not None
              else None,
          )
          key, value = torch.chunk(kv, chunks=2, dim=-1)
  
      query = query * self.scaling
  
      # [N, S, C] --> [N, S, c] x h, where C = c * h
      query = torch.chunk(query, chunks=self.num_heads, dim=-1)
  
      # [N, T, C] --> [N, T, c] x h, where C = c * h
      value = torch.chunk(value, chunks=self.num_heads, dim=-1)
      # [N, T, C] --> [N, T, c] x h, where C = c * h
      key = torch.chunk(key, chunks=self.num_heads, dim=-1)
  
      wt_out = []
      for h in range(self.num_heads):
          attn_h = torch.matmul(query[h], key[h].transpose(-1, -2))
          attn_h = self.softmax(attn_h)
          attn_h = self.attn_dropout(attn_h)
          out_h = torch.matmul(attn_h, value[h])
          wt_out.append(out_h)
  
      wt_out = torch.cat(wt_out, dim=-1)
      wt_out = self.out_proj(wt_out)
      return wt_out
  
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 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward_tracing(x_q, x_kv, key_padding_mask, attn_mask) and returns: wt_out