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 forward(self, x_q: Tensor) -> Tensor:
      # [N, P, C]
      b_sz, n_patches, in_channels = x_q.shape
  
      # self-attention
      # [N, P, C] -> [N, P, 3C] -> [N, P, 3, h, c] where C = hc
      qkv = self.qkv_proj(x_q).reshape(
          b_sz, n_patches, 3, self.num_heads, -1
      )
  
      # [N, P, 3, h, c] -> [N, h, 3, P, C]
      qkv = qkv.transpose(1, 3).contiguous()
  
      # [N, h, 3, P, C] -> [N, h, P, C] x 3
      query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
  
      query = query * self.scaling
  
      # [N h, P, c] -> [N, h, c, P]
      key = key.transpose(-1, -2)
  
      # QK^T
      # [N, h, P, c] x [N, h, c, P] -> [N, h, P, P]
      attn = torch.matmul(query, key)
      attn = self.softmax(attn)
      attn = self.attn_dropout(attn)
  
      # weighted sum
      # [N, h, P, P] x [N, h, P, c] -> [N, h, P, c]
      out = torch.matmul(attn, value)
  
      # [N, h, P, c] -> [N, P, h, c] -> [N, P, C]
      out = out.transpose(1, 2).reshape(b_sz, n_patches, -1)
      out = self.out_proj(out)
  
      return out
  
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 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(x_q) and returns: out