class_docstr: The Tensor Parallel MLP prompt encoder network that is used to generate the virtual 
  token embeddings for p-tuning. It only have two layers.
docstr: Initializes the Tensor Model parallel MLP PromptEncoderMLP module.
  Args:
      config: the model parallel config used my megatron core
      total_virtual_tokens: the total number of vitural tokens
      hidden_size: hidden dimension
      output_size:  the output dimension
      init_std: the MLP init std value 
docstr_tok: ["Initializes","the","Tensor","Model","parallel","MLP","PromptEncoderMLP","module",".","Args",":","config",":","the","model","parallel","config","used","my","megatron","core","total_virtual_tokens",":","the","total","number","of","vitural","tokens","hidden_size",":","hidden","dimension","output_size",":","the","output","dimension","init_std",":","the","MLP","init","std","value"]
code: def __init__(
      self,
      config: ModelParallelConfig,
      total_virtual_tokens: int,
      hidden_size: int,
      output_size: int,
      init_std: float,
  ):
      """
      Initializes the Tensor Model parallel MLP PromptEncoderMLP module.
      Args:
          config: the model parallel config used my megatron core
          total_virtual_tokens: the total number of vitural tokens
          hidden_size: hidden dimension
          output_size:  the output dimension
          init_std: the MLP init std value
      """
      super().__init__()
      self.hidden_size = hidden_size
      self.output_size = output_size
      self.total_virtual_tokens = total_virtual_tokens
      self.activation = "gelu"
  
      config = copy.deepcopy(config)
      config.sequence_parallel = False
      config.gradient_accumulation_fusion = False
  
      self.first = tensor_parallel.ColumnParallelLinear(
          self.output_size,
          self.hidden_size,
          config=config,
          gather_output=False,
          init_method=init_method_normal(init_std),
          skip_bias_add=True,
          bias=True,
      )
      self.second = tensor_parallel.RowParallelLinear(
          self.hidden_size,
          self.output_size,
          config=config,
          input_is_parallel=True,
          init_method=init_method_normal(init_std),
          skip_bias_add=True,
          bias=True,
      )
  
imports: import copy
  import enum
  from typing import Dict, Optional
  import torch
  import torch.nn.init
  from torch import nn
  from nemo.collections.nlp.modules.common.megatron.fused_bias_gelu import fused_bias_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, init_method_normal
  from nemo.core.classes import Exportable, NeuralModule
  from nemo.core.classes.common import typecheck
label_desc: Use image node_id 1 to create a new TPMLP object from inherited base classes: NeuralModule, Exportable with example: obj = TPMLP(config, total_virtual_tokens, hidden_size, output_size, init_std)