class_docstr: This class defines the pre-norm transformer encoder with linear self-attention in `MobileViTv2 <https://arxiv.org/abs/2206.02680>`_ paper
  Args:
      opts: command line arguments
      embed_dim (int): :math:`C_{in}` from an expected input of size :math:`(B, C_{in}, P, N)`
      ffn_latent_dim (int): Inner dimension of the FFN
      attn_dropout (Optional[float]): Dropout rate for attention in multi-head attention. Default: 0.0
      dropout (Optional[float]): Dropout rate. Default: 0.0
      ffn_dropout (Optional[float]): Dropout between FFN layers. Default: 0.0
      norm_layer (Optional[str]): Normalization layer. Default: layer_norm_2d
  
  Shape:
      - Input: :math:`(B, C_{in}, P, N)` where :math:`B` is batch size, :math:`C_{in}` is input embedding dim,
          :math:`P` is number of pixels in a patch, and :math:`N` is number of patches,
      - Output: same shape as the input
class_docstr_tok: ["This","class","defines","the","pre-norm","transformer","encoder","with","linear","self-attention","in","`","MobileViTv2","<","https",":","\/\/arxiv.org\/abs\/2206.02680",">","`","_","paper","Args",":","opts",":","command","line","arguments","embed_dim","(","int",")",":",":","math",":","`","C_","{","in","}","`","from","an","expected","input","of","size",":","math",":","`","(","B",",","C_","{","in","}",",","P",",","N",")","`","ffn_latent_dim","(","int",")",":","Inner","dimension","of","the","FFN","attn_dropout","(","Optional","[","float","]",")",":","Dropout","rate","for","attention","in","multi-head","attention",".","Default",":","0.0","dropout","(","Optional","[","float","]",")",":","Dropout","rate",".","Default",":","0.0","ffn_dropout","(","Optional","[","float","]",")",":","Dropout","between","FFN","layers",".","Default",":","0.0","norm_layer","(","Optional","[","str","]",")",":","Normalization","layer",".","Default",":","layer_norm_2d","Shape",":","-","Input",":",":","math",":","`","(","B",",","C_","{","in","}",",","P",",","N",")","`","where",":","math",":","`","B","`","is","batch","size",",",":","math",":","`","C_","{","in","}","`","is","input","embedding","dim",",",":","math",":","`","P","`","is","number","of","pixels","in","a","patch",",","and",":","math",":","`","N","`","is","number","of","patches",",","-","Output",":","same","shape","as","the","input"]
code: def __init__(
      self,
      opts,
      embed_dim: int,
      ffn_latent_dim: int,
      attn_dropout: Optional[float] = 0.0,
      dropout: Optional[float] = 0.1,
      ffn_dropout: Optional[float] = 0.0,
      norm_layer: Optional[str] = "layer_norm_2d",
      *args,
      **kwargs,
  ) -> None:
      super().__init__()
      attn_unit = LinearSelfAttention(
          opts,
          embed_dim=embed_dim,
          attn_dropout=attn_dropout,
          bias=True,
      )
  
      self.pre_norm_attn = nn.Sequential(
          get_normalization_layer(
              opts=opts, norm_type=norm_layer, num_features=embed_dim
          ),
          attn_unit,
          Dropout(p=dropout),
      )
  
      self.pre_norm_ffn = nn.Sequential(
          get_normalization_layer(
              opts=opts, norm_type=norm_layer, num_features=embed_dim
          ),
          ConvLayer2d(
              opts=opts,
              in_channels=embed_dim,
              out_channels=ffn_latent_dim,
              kernel_size=1,
              stride=1,
              bias=True,
              use_norm=False,
              use_act=True,
          ),
          Dropout(p=ffn_dropout),
          ConvLayer2d(
              opts=opts,
              in_channels=ffn_latent_dim,
              out_channels=embed_dim,
              kernel_size=1,
              stride=1,
              bias=True,
              use_norm=False,
              use_act=False,
          ),
          Dropout(p=dropout),
      )
  
      self.embed_dim = embed_dim
      self.ffn_dim = ffn_latent_dim
      self.ffn_dropout = ffn_dropout
      self.std_dropout = dropout
      self.attn_fn_name = attn_unit.__repr__()
      self.norm_name = norm_layer
  
file_path: ml-cvnets/cvnets/modules/transformer.py
imports: import argparse
  from typing import Optional
  from torch import Tensor, nn
  from cvnets.layers import ConvLayer2d, Dropout, Identity, LinearLayer, LinearSelfAttention, MultiHeadAttention, SingleHeadAttention, StochasticDepth, get_normalization_layer
  from cvnets.layers.activation import build_activation_layer
  from cvnets.modules import BaseModule
  from utils import logger
label_desc: Use image node_id 1 to create a new LinearAttnFFN object from inherited base classes: BaseModule with example: obj = LinearAttnFFN(opts, embed_dim, ffn_latent_dim, attn_dropout, dropout, ffn_dropout, norm_layer)