class_docstr: This class defines the pre-norm `Transformer encoder <https://arxiv.org/abs/1706.03762>`_
  Args:
      embed_dim (int): :math:`C_{in}` from an expected input of size :math:`(N, P, C_{in})`
      ffn_latent_dim (int): Inner dimension of the FFN
      num_heads (int) : Number of heads in multi-head attention. Default: 8
      attn_dropout (float): Dropout rate for attention in multi-head attention. Default: 0.0
      dropout (float): Dropout rate. Default: 0.0
      ffn_dropout (float): Dropout between FFN layers. Default: 0.0
  
  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","class","defines","the","pre-norm","`","Transformer","encoder","<","https",":","\/\/arxiv.org\/abs\/1706.03762",">","`","_","Args",":","embed_dim","(","int",")",":",":","math",":","`","C_","{","in","}","`","from","an","expected","input","of","size",":","math",":","`","(","N",",","P",",","C_","{","in","}",")","`","ffn_latent_dim","(","int",")",":","Inner","dimension","of","the","FFN","num_heads","(","int",")",":","Number","of","heads","in","multi-head","attention",".","Default",":","8","attn_dropout","(","float",")",":","Dropout","rate","for","attention","in","multi-head","attention",".","Default",":","0.0","dropout","(","float",")",":","Dropout","rate",".","Default",":","0.0","ffn_dropout","(","float",")",":","Dropout","between","FFN","layers",".","Default",":","0.0","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,
      ffn_latent_dim: int,
      num_heads: Optional[int] = 8,
      attn_dropout: Optional[float] = 0.0,
      dropout: Optional[float] = 0.0,
      ffn_dropout: Optional[float] = 0.0,
      *args,
      **kwargs,
  ) -> None:
      super().__init__()
  
      attn_unit = MultiHeadAttention(
          embed_dim, num_heads, attn_dropout=attn_dropout, bias=True
      )
  
      self.pre_norm_mha = nn.Sequential(
          nn.LayerNorm(embed_dim), attn_unit, nn.Dropout(p=dropout)
      )
  
      self.pre_norm_ffn = nn.Sequential(
          nn.LayerNorm(embed_dim),
          nn.Linear(
              in_features=embed_dim,
              out_features=ffn_latent_dim,
              bias=True,
          ),
          nn.SiLU(),
          nn.Dropout(p=ffn_dropout),
          nn.Linear(
              in_features=ffn_latent_dim,
              out_features=embed_dim,
              bias=True,
          ),
          nn.Dropout(p=dropout),
      )
      self.embed_dim = embed_dim
      self.ffn_dim = ffn_latent_dim
      self.ffn_dropout = ffn_dropout
      self.std_dropout = dropout
  
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 TransformerEncoder object from inherited base classes: nn with example: obj = TransformerEncoder(embed_dim, ffn_latent_dim, num_heads, attn_dropout, dropout, ffn_dropout)