class_docstr: T5 Language model.
  
  Args:
      config (TransformerConfig): transformer config
  
      transformer_encoder_layer_spec (ModuleSpec): transformer layer customization specs for encoder
  
      transformer_decoder_layer_spec (ModuleSpec): transformer layer customization specs for decoder
              
      vocab_size (int): vocabulary size
  
      max_sequence_length (int): maximum size of sequence. This is used for positional embedding
  
      pre_process (bool): Include embedding layer (used with pipeline parallelism)
      post_process (bool): Include an output layer (used with pipeline parallelism)
  
      fp16_lm_cross_entropy (bool, optional): Defaults to False
  
      parallel_output (bool): Do not gather the outputs, keep them split across tensor parallel ranks
  
      share_embeddings_and_output_weights (bool): When True, input embeddings and output logit weights are
          shared. Defaults to False.
  
      position_embedding_type (string): Position embedding type. Options ['learned_absolute', 'rope'].
          Defaults is 'learned_absolute'.
  
      rotary_percent (float): Percent of rotary dimension to use for rotary position embeddings.
          Defaults to 1.0 (100%). Ignored unless position_embedding_type is 'rope'.
  
      seq_len_interpolation_factor (float): scale of linearly interpolating RoPE for longer sequences.
          The value must be a float larger than 1.0. Defaults to None.
class_docstr_tok: ["T5","Language","model",".","Args",":","config","(","TransformerConfig",")",":","transformer","config","transformer_encoder_layer_spec","(","ModuleSpec",")",":","transformer","layer","customization","specs","for","encoder","transformer_decoder_layer_spec","(","ModuleSpec",")",":","transformer","layer","customization","specs","for","decoder","vocab_size","(","int",")",":","vocabulary","size","max_sequence_length","(","int",")",":","maximum","size","of","sequence",".","This","is","used","for","positional","embedding","pre_process","(","bool",")",":","Include","embedding","layer","(","used","with","pipeline","parallelism",")","post_process","(","bool",")",":","Include","an","output","layer","(","used","with","pipeline","parallelism",")","fp16_lm_cross_entropy","(","bool",",","optional",")",":","Defaults","to","False","parallel_output","(","bool",")",":","Do","not","gather","the","outputs",",","keep","them","split","across","tensor","parallel","ranks","share_embeddings_and_output_weights","(","bool",")",":","When","True",",","input","embeddings","and","output","logit","weights","are","shared",".","Defaults","to","False",".","position_embedding_type","(","string",")",":","Position","embedding","type",".","Options","[","'learned_absolute","'",",","'rope","'","]",".","Defaults","is","'learned_absolute","'",".","rotary_percent","(","float",")",":","Percent","of","rotary","dimension","to","use","for","rotary","position","embeddings",".","Defaults","to","1.0","(","100","%",")",".","Ignored","unless","position_embedding_type","is","'rope","'",".","seq_len_interpolation_factor","(","float",")",":","scale","of","linearly","interpolating","RoPE","for","longer","sequences",".","The","value","must","be","a","float","larger","than","1.0",".","Defaults","to","None","."]
docstr: Forward pass.
  
  Args:
      encoder_input_ids (Tensor): input ids for encoder
      decoder_input_ids (Tensor): input ids for decoder
      encoder_attn_mask (Tensor): self-attention mask for encoder
      decoder_attn_mask (Tensor): self-attention mask for decoder
      encoder_decoder_attn_mask (Tensor): cross-attention mask between encoder and decoder
      lm_labels (Tensor): labels for decoder output
      inference_params (InferenceParams): relevant arguments for inferencing
  
  Returns:
      Tensor: loss tensor
docstr_tok: ["Forward","pass",".","Args",":","encoder_input_ids","(","Tensor",")",":","input","ids","for","encoder","decoder_input_ids","(","Tensor",")",":","input","ids","for","decoder","encoder_attn_mask","(","Tensor",")",":","self-attention","mask","for","encoder","decoder_attn_mask","(","Tensor",")",":","self-attention","mask","for","decoder","encoder_decoder_attn_mask","(","Tensor",")",":","cross-attention","mask","between","encoder","and","decoder","lm_labels","(","Tensor",")",":","labels","for","decoder","output","inference_params","(","InferenceParams",")",":","relevant","arguments","for","inferencing","Returns",":","Tensor",":","loss","tensor"]
code: def forward(
      self,
      encoder_input_ids: Tensor,
      decoder_input_ids: Tensor,
      encoder_attn_mask: Tensor,
      decoder_attn_mask: Tensor,
      encoder_decoder_attn_mask: Tensor,
      lm_labels: Tensor = None,
      inference_params: InferenceParams = None,
  ) -> Tensor:
      """Forward pass.
  
      Args:
          encoder_input_ids (Tensor): input ids for encoder
          decoder_input_ids (Tensor): input ids for decoder
          encoder_attn_mask (Tensor): self-attention mask for encoder
          decoder_attn_mask (Tensor): self-attention mask for decoder
          encoder_decoder_attn_mask (Tensor): cross-attention mask between encoder and decoder
          lm_labels (Tensor): labels for decoder output
          inference_params (InferenceParams): relevant arguments for inferencing
  
      Returns:
          Tensor: loss tensor
      """
  
      (
          encoder_attn_mask,
          decoder_attn_mask,
          encoder_decoder_attn_mask,
      ) = t5_extended_attention_mask(
          [
              encoder_attn_mask,
              decoder_attn_mask,
              encoder_decoder_attn_mask,
          ]
      )
      encoder_position_ids = t5_position_ids(encoder_input_ids)
      decoder_position_ids = t5_position_ids(decoder_input_ids)
  
      ## Encoder forward
      # Encoder embedding.
      if self.pre_process:
          encoder_input = self.embedding(
              input_ids=encoder_input_ids,
              position_ids=encoder_position_ids,
          )
      else:
          # intermediate stage of pipeline
          encoder_input = None
  
      # Rotary positional embeddings
      rotary_pos_emb = None
      if self.position_embedding_type == "rope":
          rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
              inference_params, self.encoder, encoder_input, self.config
          )
          rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
  
      # Run encoder.
      encoder_hidden_states = self.encoder(
          hidden_states=encoder_input,
          attention_mask=encoder_attn_mask,
          inference_params=inference_params,
          rotary_pos_emb=rotary_pos_emb,
      )
  
      ## Decoder forward
      # Decoder embedding.
      if self.pre_process:
          decoder_input = self.embedding(
              input_ids=decoder_input_ids,
              position_ids=decoder_position_ids,
          )
      else:
          # intermediate stage of pipeline
          decoder_input = None  ### should it take encoder_hidden_states
  
      # Rotary positional embeddings
      rotary_pos_emb = None
      if self.position_embedding_type == "rope":
          rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
              inference_params, self.decoder, decoder_input, self.config
          )
          rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
  
      # Run decoder.
      decoder_hidden_states = self.decoder(
          hidden_states=decoder_input,
          attention_mask=decoder_attn_mask,
          context=encoder_hidden_states,
          context_mask=encoder_decoder_attn_mask,
          inference_params=inference_params,
          rotary_pos_emb=rotary_pos_emb,
      )
  
      # Return if not post_process
      if not self.post_process:
          return decoder_hidden_states
  
      # logits and loss
      output_weight = None
      if self.share_embeddings_and_output_weights:
          output_weight = self.shared_embedding_or_output_weight()
      logits = self.lm_head(
          decoder_hidden_states, word_embeddings_weight=output_weight
      )
  
      if lm_labels is None:
          # [s b h] => [b s h]
          return logits.transpose(0, 1).contiguous()
  
      loss = self.compute_language_model_loss(lm_labels, logits)
  
      return loss
  
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 T5Model obj's underlying member method code with example usage: obj.forward(encoder_input_ids, decoder_input_ids, encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask, lm_labels, inference_params) and returns: loss, decoder_hidden_states, logits