class_docstr: Megatron encoder-decoder base class
class_docstr_tok: ["Megatron","encoder-decoder","base","class"]
docstr: tokens_enc - encoder input tokens
  enc_mask - corresponding mask
  encoder_input - encoder input (bypass tokens), if given tokens_enc can be None.
  batch_data - passed directly to all hidden transformations and losses. 
               Can be used to pass additional data like class label.
               Format is not defined and should match the expected format of the used hiddens modules.
docstr_tok: ["tokens_enc","-","encoder","input","tokens","enc_mask","-","corresponding","mask","encoder_input","-","encoder","input","(","bypass","tokens",")",",","if","given","tokens_enc","can","be","None",".","batch_data","-","passed","directly","to","all","hidden","transformations","and","losses",".","Can","be","used","to","pass","additional","data","like","class","label",".","Format","is","not","defined","and","should","match","the","expected","format","of","the","used","hiddens","modules","."]
code: def encode(
      self,
      tokens_enc,
      enc_mask,
      encoder_input=None,
      batch_data=None,
      reconfigure_microbatch=True,
  ):
      """
      tokens_enc - encoder input tokens
      enc_mask - corresponding mask
      encoder_input - encoder input (bypass tokens), if given tokens_enc can be None.
      batch_data - passed directly to all hidden transformations and losses.
                   Can be used to pass additional data like class label.
                   Format is not defined and should match the expected format of the used hiddens modules.
      """
      # Check whether the DDP is initialized. This is needed when running inference outside of training loop.
      if parallel_state.is_unitialized():
  
          def dummy():
              return
  
          if self.trainer.strategy.launcher is not None:
              self.trainer.strategy.launcher.launch(
                  dummy, trainer=self.trainer
              )
          self.trainer.strategy.setup_environment()
  
          # Reconfigure microbatch sizes here because on model restore, this will contain the micro/global batch configuration used while training.
          if reconfigure_microbatch:
              _reconfigure_microbatch_calculator(
                  rank=0,  # This doesn't matter since it is only used for logging
                  rampup_batch_size=None,
                  global_batch_size=1,
                  micro_batch_size=1,  # Make sure that there is no "grad acc" while decoding.
                  data_parallel_size=1,  # We check above to make sure that dataparallel size is always 1 at inference.
              )
  
      # If classes that inherit from this class are using a different tokenizer,
      app_state = AppState()
      if tokens_enc is not None:
          global_batch_per_gpu = tokens_enc.size(0)
          encoder_seq_length = tokens_enc.size(1)
      else:
          global_batch_per_gpu = encoder_input.size(1)
          encoder_seq_length = encoder_input.size(0)
  
      num_micro_batches_before_decode = get_num_microbatches()
      # Reconfigure microbatch calculator here to set num microbatches to 1 while decoding since its not clear how to decode with "grad acc".
      # reconfigure back to how things were before encode
      if reconfigure_microbatch:
          _reconfigure_microbatch_calculator(
              rank=app_state.global_rank,
              rampup_batch_size=None,
              global_batch_size=global_batch_per_gpu
              * parallel_state.get_data_parallel_world_size(),
              micro_batch_size=global_batch_per_gpu,  # Make sure that there is no "grad acc" while decoding.
              data_parallel_size=parallel_state.get_data_parallel_world_size(),
          )
      tensor_shape = [
          encoder_seq_length,
          global_batch_per_gpu,
          self.cfg.encoder.hidden_size,
      ]
  
      # build input arguments description
      if tokens_enc is not None:
          batch_for_pipeline = [tokens_enc, enc_mask, batch_data]
          arg_names = ["enc_input_ids", "enc_attn_mask", "batch_data"]
      else:
          if encoder_input is None:
              raise ValueError(
                  "At least one of tokens_enc and encoder_input must be provided with not None value"
              )
  
          batch_for_pipeline = [enc_mask]
          arg_names = ["enc_attn_mask"]
  
      if encoder_input is not None:
          batch_for_pipeline.append(encoder_input)
          arg_names.append("enc_input")
  
      forward_step_func = self._get_forward_output_only_func(
          arg_names=arg_names,
          output_name="hiddens",
          output_enc_hidden_only=True,
      )
  
      fwd_bwd_func = get_forward_backward_func()
  
      # Counter intuitively, we need to set decoder_sequence_length=encoder_seq_length
      # because while running `.encode()`, the last hidden states from encoder are passed through
      # as identity through the pipeline.
      # Setting it to anything else will cause hanging due to tensor shape mismatches.
      output_tensor = fwd_bwd_func(
          forward_step_func=forward_step_func,
          data_iterator=iter(
              [
                  batch_for_pipeline,
              ]
          ),
          model=[self.enc_dec_model],
          forward_only=True,
          num_microbatches=1,
          seq_length=encoder_seq_length,
          decoder_seq_length=encoder_seq_length,
          micro_batch_size=get_micro_batch_size(),
      )
  
      if output_tensor:
          output_tensor = output_tensor[0]["hiddens"]
      else:
          output_tensor = torch.zeros(
              tensor_shape, dtype=self.autocast_dtype
          ).cuda()
  
      if self.cfg.get("pipeline_model_parallel_size", 1) > 1:
          # Broadcast from the last pipeline stage to all other model-parallel ranks.
          torch.distributed.broadcast(
              output_tensor,
              parallel_state.get_pipeline_model_parallel_last_rank(),
              group=parallel_state.get_pipeline_model_parallel_group(),
          )
  
      # Reset microbatch calculator to what it was before decoding.
      if reconfigure_microbatch:
          _reconfigure_microbatch_calculator(
              rank=app_state.global_rank,
              rampup_batch_size=None,
              global_batch_size=global_batch_per_gpu
              * parallel_state.get_data_parallel_world_size(),
              micro_batch_size=global_batch_per_gpu
              // num_micro_batches_before_decode,
              data_parallel_size=parallel_state.get_data_parallel_world_size(),
          )
  
      # Return the output tensor of encoder and transpose from [seq_len, batch, hidden] to [batch, seq_len, hidden]
      return output_tensor.transpose(1, 0)
  
imports: import copy
  import functools
  import inspect
  from typing import Any, Dict, List, Optional
  import torch
  from omegaconf import OmegaConf, open_dict
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.accelerators import CPUAccelerator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.build_model import build_model
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.token_level_encoder_decoder import MegatronTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_utils import compute_beam_search_len_penalty, get_sampling_token_fn
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 38 for calling the MegatronLMEncoderDecoderModel obj's underlying member method code with example usage: obj.encode(tokens_enc, enc_mask, encoder_input, batch_data, reconfigure_microbatch) and returns: output_tensor