code: def compute_splits(
      self,
      model,
      partitions,
      idx,
      tp_rank,
      pp_rank,
      pp_split_rank,
      tp_size,
      pp_size,
  ):
      splits = []
  
      # Backup index when EncDec models reset the index to fill in the first embedding matrices (when pp split rank == pp rank)
      computed_index = idx
  
      # This is the PP X TP Y model with partial parameters present in correct order.
      # We need to extract the parameters from the global map in reverse order to fill in the
      # parameters of this model in forward order.
      for param_name, param in model.named_parameters():
          # Since we are moving forward, we may reach the end of the global map
          # but T5 has an additional word embedding as its first two parameter when pp split rank == pp rank
          # Therefore we check for this, and update the index to the parameter of the PP 0 TP 0 rank
          # which holds the parameters of the embedding.
          if self.enc_dec_share_token_embeddings_count:
              logging.info(
                  "EncDec models decoder shares embedding with encoder, resetting index"
              )
              idx = (
                  2 - self.enc_dec_share_token_embeddings_count
              )  # 0th index is vocab embedding, 1 is pos embedding, 2 is embedding count
  
          # Since we are moving forward, we may reach the end of the global map
          # but T5 has an additional word embedding as randomly located in the decoder when
          # when pp rank > pp_split_rank.
          # Therefore we check for this, and skip the parameter of the current TP X PP Y module
          # and fill this parameter later.
          if (
              self.shared_enc_dec_embeddings_intermediate
              and param_name == "enc_dec_model.word_embeddings.weight"
          ):
              logging.info(
                  "EncDec models decoder shares embedding with encoder in intermediate pos, skipping module for later update"
              )
              continue
  
          debug_log_split_param_diff(idx, param, param_name, partitions)
  
          # Tensor Parallel Splitting
          split = compute_tp_splits(
              param_name,
              param,
              partitions,
              idx,
              tp_size,
              pp_size,
              pp_rank,
              pp_split_rank,
              self.megatron_legacy,
              model.cfg,
          )
  
          splits.append(split)
          idx += 1
  
          # When pp split rank is equal to current pp rank, we need to first inject the encoder embeddings
          # and then reset the index to the originally computed index
          if self.enc_dec_share_token_embeddings_count > 0:
              if self.enc_dec_share_token_embeddings_count - 1 == 0:
                  idx = computed_index
  
              self.enc_dec_share_token_embeddings_count -= 1
  
      # Inject the EncDec shared embeddings intermediate tensor
      # at one random location in the decoder of this TP PP rank.
      # Note that usually it is the second last tensor, but to avoid specific index we search for it
      # again.
      if self.shared_enc_dec_embeddings_intermediate:
          for param_name, param in model.named_parameters():
              if param_name == "enc_dec_model.word_embeddings.weight":
                  logging.info(
                      "Found intermediate shared embedding, injecting"
                  )
                  split = compute_tp_splits(
                      param_name,
                      param,
                      partitions,
                      global_idx=0,
                      tp_size=tp_size,
                      pp_size=pp_size,
                      pp_rank=pp_rank,
                      pp_split_rank=pp_split_rank,
                      megatron_legacy=self.megatron_legacy,
                      model_cfg=model.cfg,
                  )
                  splits.insert(
                      self.intermediate_shared_embedding_location, split
                  )
                  break
  
      return idx, splits
  
imports: import os
  import shutil
  import tarfile
  import tempfile
  from argparse import ArgumentParser
  from typing import Dict, List
  import torch
  import torch.nn
  from omegaconf import OmegaConf, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.nlp.parts.nlp_overrides import NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE, GradScaler, MegatronHalfPrecisionPlugin, NLPDDPStrategy, NLPSaveRestoreConnector, PipelineMixedPrecisionPlugin
  from nemo.utils import logging, model_utils
  from nemo.utils.app_state import AppState
label_desc: Use image node_id 3 for calling the T5Handler obj's underlying member method code with example usage: obj.compute_splits(model, partitions, idx, tp_rank, pp_rank, pp_split_rank, tp_size, pp_size) and returns: idx,  splits