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, P, C_{in})`
      num_heads (int): Number of heads in multi-head attention
      attn_dropout (float): Attention dropout. Default: 0.0
      bias (bool): Use bias or not. Default: ``True``
  
  Shape:
      - Input: :math:`(N, P, C_{in})` where :math:`N` is batch size, :math:`P` is number of patches,
      and :math:`C_{in}` is input embedding dim
      - Output: same shape as the input
class_docstr_tok: ["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",",","P",",","C_","{","in","}",")","`","num_heads","(","int",")",":","Number","of","heads","in","multi-head","attention","attn_dropout","(","float",")",":","Attention","dropout",".","Default",":","0.0","bias","(","bool",")",":","Use","bias","or","not",".","Default",":","``","True","``","Shape",":","-","Input",":",":","math",":","`","(","N",",","P",",","C_","{","in","}",")","`","where",":","math",":","`","N","`","is","batch","size",",",":","math",":","`","P","`","is","number","of","patches",",","and",":","math",":","`","C_","{","in","}","`","is","input","embedding","dim","-","Output",":","same","shape","as","the","input"]
code: def __init__(
      self,
      embed_dim: int,
      num_heads: int,
      attn_dropout: float = 0.0,
      bias: bool = True,
      *args,
      **kwargs,
  ) -> None:
      super().__init__()
      if embed_dim % num_heads != 0:
          raise ValueError(
              "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 = nn.Linear(
          in_features=embed_dim, out_features=3 * embed_dim, bias=bias
      )
  
      self.attn_dropout = nn.Dropout(p=attn_dropout)
      self.out_proj = nn.Linear(
          in_features=embed_dim, out_features=embed_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
  
file_path: deep-learning-for-image-processing/pytorch_classification/MobileViT/transformer.py
imports: from typing import Optional
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(embed_dim, num_heads, attn_dropout, bias)