class_docstr: The prompt encoder network that is used to generate the virtual 
  token embeddings for p-tuning.
docstr: Initializes the PromptEncoder module.
  Args:
      config: the model parallel config used my megatron core
      total_virtual_tokens: the total number of vitural tokens
      hidden_size: hidden dimension
      lstm_dropout: the dropout used for the LSTM
      num_layers: number of layers used in the LSTM
      init_std: used for TPMLP encoder type to initialize the mlp weights
docstr_tok: ["Initializes","the","PromptEncoder","module",".","Args",":","config",":","the","model","parallel","config","used","my","megatron","core","total_virtual_tokens",":","the","total","number","of","vitural","tokens","hidden_size",":","hidden","dimension","lstm_dropout",":","the","dropout","used","for","the","LSTM","num_layers",":","number","of","layers","used","in","the","LSTM","init_std",":","used","for","TPMLP","encoder","type","to","initialize","the","mlp","weights"]
code: def __init__(
      self,
      config: ModelParallelConfig,
      encoder_type: enum,
      total_virtual_tokens: int,
      token_dim: int,
      hidden_size,
      lstm_dropout: float,
      num_layers: int,
      init_std: float,
      taskname: str = "taskname",
  ):
      """
      Initializes the PromptEncoder module.
      Args:
          config: the model parallel config used my megatron core
          total_virtual_tokens: the total number of vitural tokens
          hidden_size: hidden dimension
          lstm_dropout: the dropout used for the LSTM
          num_layers: number of layers used in the LSTM
          init_std: used for TPMLP encoder type to initialize the mlp weights
      """
      super().__init__()
      self.token_dim = token_dim
      self.input_size = token_dim
      self.output_size = token_dim
      self.hidden_size = hidden_size
      self.total_virtual_tokens = total_virtual_tokens
      self.encoder_type = encoder_type
      self.activation = "gelu"
      self.init_std = init_std
      self.taskname = taskname
  
      # Set fixed indicies for forward pass
      self.register_buffer(
          "indices",
          torch.LongTensor(list(range(self.total_virtual_tokens))),
      )
  
      # embedding
      self.embedding = torch.nn.Embedding(
          self.total_virtual_tokens, self.token_dim
      )
      self.inference_table = InferenceTable(
          taskname, self.token_dim, self.total_virtual_tokens
      )
  
      if self.encoder_type == PromptEncoderType.EMBEDDING:
          init.xavier_normal_(self.embedding.weight)
      elif self.encoder_type == PromptEncoderType.LSTM:
          # LSTM
          self.lstm_head = torch.nn.LSTM(
              input_size=self.input_size,
              hidden_size=self.hidden_size,
              num_layers=num_layers,
              dropout=lstm_dropout,
              bidirectional=True,
              batch_first=True,
          )
  
          self.mlp_head = nn.Sequential(
              nn.Linear(self.hidden_size * 2, self.hidden_size * 2),
              nn.ReLU(),
              nn.Linear(self.hidden_size * 2, self.output_size),
          )
  
      elif self.encoder_type == PromptEncoderType.MLP:
          if num_layers <= 1:
              raise ValueError(
                  "The MLP prompt encoder must have at least 2 layers, and exactly 2 layers is recommended."
              )
  
          layers = [
              nn.Linear(self.input_size, self.hidden_size),
              nn.ReLU(),
          ]
          for _ in range(num_layers - 2):
              layers.extend(
                  [
                      nn.Linear(self.hidden_size, self.hidden_size),
                      nn.ReLU(),
                  ]
              )
  
          layers.append(nn.Linear(self.hidden_size, self.output_size))
          self.mlp_head = nn.Sequential(*layers)
  
      elif self.encoder_type == PromptEncoderType.TPMLP:
          self.tpmlp = TPMLP(
              config,
              self.total_virtual_tokens,
              self.hidden_size,
              self.output_size,
              self.init_std,
          )
      else:
          raise ValueError(
              "Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM."
          )
  
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 PromptEncoder object from inherited base classes: NeuralModule, Exportable with example: obj = PromptEncoder(config, encoder_type, total_virtual_tokens, token_dim, hidden_size, lstm_dropout, num_layers, init_std, taskname)