class_docstr: Token-based (input/output is tokens) retrieval encoder-decoder model
class_docstr_tok: ["Token-based","(","input\/output","is","tokens",")","retrieval","encoder-decoder","model"]
docstr: Return value is per token / per dimension (i.e., non collapsed loss value)
docstr_tok: ["Return","value","is","per","token","\/","per","dimension","(","i.e.",",","non","collapsed","loss","value",")"]
code: def forward(
      self,
      input_ids,
      input_attn_mask,
      retrieved_ids,
      retrieved_attn_mask,
      token_type_ids=None,
      labels=None,
      input_emb=None,
      set_inference_key_value_memory=False,
      inference_max_sequence_len=None,
      neighbors=None,
      position_ids=None,
  ):
      """
      Return value is per token / per dimension (i.e., non collapsed loss value)
      """
      eod_positions = None
      retrieved_emb = None
      if (
          input_ids is not None
          and self.eod_id is not None
          and not self.megatron_lm_compatible
      ):
          # do not reset attention for megatron lm compatible model
          eod_positions = torch.where(input_ids == self.eod_id)
  
      if input_emb is None:
          if self.pre_process and self.add_encoder:
              # encoder embeddings
              if self.add_abs_position_embedding:
                  input_position_ids = position_ids
              else:
                  input_position_ids = None
              input_emb = self.encoder_embedding(
                  input_ids,
                  input_position_ids,
                  token_type_ids=token_type_ids,
              )
          else:
              input_emb = None
  
      if retrieved_ids is not None:
          if self.add_abs_position_embedding:
              seq_length = retrieved_ids.size(-1)
              retrieved_position_ids = torch.arange(
                  seq_length,
                  dtype=torch.long,
                  device=retrieved_ids.device,
              )
              retrieved_position_ids = (
                  retrieved_position_ids.unsqueeze(0)
                  .expand_as(retrieved_ids)
                  .clone()
              )
          else:
              retrieved_position_ids = None
          retrieved_emb = self.encoder_embedding(
              retrieved_ids, retrieved_position_ids
          )
  
      if self.add_decoder:
          hidden = self.pre_decoder(
              input_emb,
              input_attn_mask,
              eod_positions=eod_positions,
              set_inference_key_value_memory=set_inference_key_value_memory,
              inference_max_sequence_len=inference_max_sequence_len,
          )
          # hidden is a tuple, (layernorm_input, layernorm_output)
          self.post_decoder.set_input_tensor(hidden)
          encoder_output = hidden[1].transpose(0, 1).contiguous()
  
      if self.add_encoder:
          if retrieved_emb is not None and neighbors is None:
              neighbors = retrieved_emb.shape[2]
          retrieved_emb = self.encoder(
              retrieved_emb,
              retrieved_attn_mask,
              context_attn_mask=input_attn_mask,
              encoder_output=encoder_output,
              set_inference_key_value_memory=set_inference_key_value_memory,
              inference_max_sequence_len=inference_max_sequence_len,
              neighbors=neighbors,
          )
  
      if self.add_decoder:
          dec_output = self.post_decoder(
              hidden,
              input_attn_mask,
              retrieved_attn_mask=retrieved_attn_mask,
              retrieved_emb=retrieved_emb,
              eod_positions=eod_positions,
              set_inference_key_value_memory=set_inference_key_value_memory,
              inference_max_sequence_len=inference_max_sequence_len,
          )
          # only transpose it for post_ln
          token_logits = self.tokens_head(
              dec_output, self.word_embeddings_weight()
          )
  
          if labels is not None:
              # [b, s] -> [s, b]
              labels = labels.transpose(0, 1).contiguous()
  
              # tensor_parallel.vocab_parallel_cross_entropy performs log_softmax and return log p(x_i|z) per token i
              if self.fp16_cross_entropy:
                  assert token_logits.dtype == torch.half
                  tokens_loss = (
                      tensor_parallel.vocab_parallel_cross_entropy(
                          token_logits, labels
                      )
                  )
              else:
                  tokens_loss = (
                      tensor_parallel.vocab_parallel_cross_entropy(
                          token_logits.float(), labels
                      )
                  )
              # [s, b] -> [b, s]
              tokens_loss = tokens_loss.transpose(0, 1).contiguous()
              return tokens_loss
          else:
              # [s, b, h] -> [b, s, h]
              token_logits = token_logits.transpose(0, 1).contiguous()
              return token_logits
  
file_path: NeMo/nemo/collections/nlp/modules/common/megatron/retrieval_token_level_encoder_decoder.py
imports: import torch
  from nemo.collections.nlp.modules.common.megatron.language_model import Embedding
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.megatron_decoders import get_decoder_model
  from nemo.collections.nlp.modules.common.megatron.megatron_encoders import get_encoder_model
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.mup.layer import MuReadout
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, init_method_normal, scaled_init_method_normal
label_desc: Use image node_id 3 for calling the MegatronRetrievalTokenLevelEncoderDecoderModule obj's underlying member method code with example usage: obj.forward(input_ids, input_attn_mask, retrieved_ids, retrieved_attn_mask, token_type_ids, labels, input_emb, set_inference_key_value_memory, inference_max_sequence_len, neighbors, position_ids) and returns: tokens_loss, token_logits