class_docstr: Masked LM head for T5
  
  Args:
      config (TransformerConfig): transformer config
      parallel_output (bool): wether output logits being distributed or not.
      vocab_size (int): vocabulary size
      pre_process (bool): Include embedding layer
      share_embeddings_and_output_weights (bool): When True, input embeddings and output logit weights are
          shared.
class_docstr_tok: ["Masked","LM","head","for","T5","Args",":","config","(","TransformerConfig",")",":","transformer","config","parallel_output","(","bool",")",":","wether","output","logits","being","distributed","or","not",".","vocab_size","(","int",")",":","vocabulary","size","pre_process","(","bool",")",":","Include","embedding","layer","share_embeddings_and_output_weights","(","bool",")",":","When","True",",","input","embeddings","and","output","logit","weights","are","shared","."]
docstr: Forward pass.
  
  Args:
      hidden_states (Tensor): output hidden states from decoder
      word_embeddings_weight (Tensor): word embedding weight
  
  Returns:
      Tensor: logits tensor
docstr_tok: ["Forward","pass",".","Args",":","hidden_states","(","Tensor",")",":","output","hidden","states","from","decoder","word_embeddings_weight","(","Tensor",")",":","word","embedding","weight","Returns",":","Tensor",":","logits","tensor"]
code: def forward(
      self, hidden_states: Tensor, word_embeddings_weight: Tensor
  ) -> Tensor:
      """Forward pass.
  
      Args:
          hidden_states (Tensor): output hidden states from decoder
          word_embeddings_weight (Tensor): word embedding weight
  
      Returns:
          Tensor: logits tensor
      """
  
      logits, _ = self.output_layer(
          hidden_states, weight=word_embeddings_weight
      )
      return logits
  
imports: import logging
  from typing import List, Literal, Optional
  import torch
  from torch import Tensor
  from megatron.core import InferenceParams, parallel_state, tensor_parallel
  from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
  from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
  from megatron.core.models.common.language_module.language_module import LanguageModule
  from megatron.core.transformer.enums import AttnMaskType, ModelType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec
  from megatron.core.transformer.transformer_block import TransformerBlock
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint
label_desc: Use image node_id 2 for calling the T5LMHead obj's underlying member method code with example usage: obj.forward(hidden_states, word_embeddings_weight) and returns: logits