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_default(
      self,
      x_q: Tensor,
      x_kv: Optional[Tensor] = None,
      key_padding_mask: Optional[Tensor] = None,
      attn_mask: Optional[Tensor] = None,
  ) -> Tensor:
      # [N, S, C]
      b_sz, S_len, in_channels = x_q.shape
  
      if x_kv is None:
          # self-attention
          # [N, S, C] --> [N, S, 3C] --> [N, S, 3, h, c] where C = hc
          qkv = self.qkv_proj(x_q).reshape(
              b_sz, S_len, 3, self.num_heads, -1
          )
          # [N, S, 3, h, c] --> [N, h, 3, S, C]
          qkv = qkv.transpose(1, 3).contiguous()
  
          # [N, h, 3, S, C] --> [N, h, S, C] x 3
          query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
      else:
          T_len = x_kv.shape[1]
  
          # cross-attention
          # [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, S, C] --> [N, S, h, c] --> [N, h, S, c]
          query = (
              query.reshape(b_sz, S_len, self.num_heads, self.head_dim)
              .transpose(1, 2)
              .contiguous()
          )
  
          # [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,
          )
          # [N, T, 2C] --> [N, T, 2, h, c]
          kv = kv.reshape(b_sz, T_len, 2, self.num_heads, self.head_dim)
          # [N, T, 2, h, c] --> [N, h, 2, T, c]
          kv = kv.transpose(1, 3).contiguous()
          key, value = kv[:, :, 0], kv[:, :, 1]
  
      query = query * self.scaling
  
      # [N h, T, c] --> [N, h, c, T]
      key = key.transpose(-1, -2)
  
      # QK^T
      # [N, h, S, c] x [N, h, c, T] --> [N, h, S, T]
      attn = torch.matmul(query, key)
  
      batch_size, num_heads, num_src_tokens, num_tgt_tokens = attn.shape
      if attn_mask is not None:
          # attn_mask shape should be the same as attn
          assert list(attn_mask.shape) == [
              batch_size,
              num_src_tokens,
              num_tgt_tokens,
          ], "Shape of attention mask should be [{}, {}, {}]. Got: {}".format(
              batch_size,
              num_src_tokens,
              num_tgt_tokens,
              attn_mask.shape,
          )
          # [N, S, T] --> [N, 1, S, T]
          attn_mask = attn_mask.unsqueeze(1)
          attn = attn + attn_mask
  
      if key_padding_mask is not None:
          # Do not attend to padding positions
          # key padding mask size is [N, T]
          assert key_padding_mask.dim() == 2 and list(
              key_padding_mask.shape
          ) == [
              batch_size,
              num_tgt_tokens,
          ], "Key_padding_mask should be 2-dimension with shape [{}, {}]. Got: {}".format(
              batch_size, num_tgt_tokens, key_padding_mask.shape
          )
          attn = attn.masked_fill(
              key_padding_mask.unsqueeze(1)
              .unsqueeze(2)
              .to(torch.bool),  # [N, T] --> [N, 1, 1, T]
              float("-inf"),
          )
  
      attn_dtype = attn.dtype
      attn_as_float = self.softmax(attn.float())
      attn = attn_as_float.to(attn_dtype)
      attn = self.attn_dropout(attn)
  
      # weighted sum
      # [N, h, S, T] x [N, h, T, c] --> [N, h, S, c]
      out = torch.matmul(attn, value)
  
      # [N, h, S, c] --> [N, S, h, c] --> [N, S, C]
      out = out.transpose(1, 2).reshape(b_sz, S_len, -1)
      out = self.out_proj(out)
  
      return 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 4 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward_default(x_q, x_kv, key_padding_mask, attn_mask) and returns: out