code: def split_tp_partition_only(
      model, partitions, tp_size, write_path=None, megatron_legacy=False
  ):
      if len(partitions) != 2:
          raise ValueError(
              "Can only split partitions of model with TP=1. For partitions of models with TP>1, merge first."
          )
  
      if tp_size < 1:
          raise ValueError("TP size must to be >= 1.")
  
      app_state = AppState()
      app_state.data_parallel_rank = 0
      app_state.pipeline_model_parallel_size = 1
      app_state.tensor_model_parallel_size = tp_size
      app_state.model_parallel_size = (
          app_state.pipeline_model_parallel_size
          * app_state.tensor_model_parallel_size
      )
  
      app_state.pipeline_model_parallel_rank = 0
      app_state.tensor_model_parallel_rank = tp_size - 1
  
      idx = 0
      splits = []
      for param_name, param in model.named_parameters():
          split = compute_tp_splits(
              param_name,
              param,
              partitions,
              idx,
              tp_size,
              pp_size=1,
              pp_rank=0,
              pp_split_rank=0,
              megatron_legacy=megatron_legacy,
              model_cfg=model.cfg,
          )
          splits.append(split)
          idx += 1
  
      # Save each of the TP ranks in reverse order
      # This is done so that the last PP rank will save the last TP rank only after all other PP TP ranks are saved
      # The final rank will then save a new NeMo file with all other ranks inside.
      write_tp_pp_split(
          model,
          splits,
          app_state,
          tp_size,
          pp_rank=0,
          write_path=write_path,
      )
  
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 10 for calling a global function with example usage: split_tp_partition_only(model, partitions, tp_size, write_path, megatron_legacy) without return types