code: def merge_tp_slices(
      ds_checkpoint, dir, slice_dir, tp_degree, name_and_shape
  ):
      name, shape = name_and_shape
      slice_base_path = os.path.join(slice_dir, name)
      param_base_path = os.path.join(dir, name)
  
      universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(
          UNIVERSAL_CHECKPOINT_INFO
      )
      replicated_parameters = universal_checkpoint_info.get(
          TP_REPLICATED_PARAMETER_PATTERNS, []
      )
      parameters_to_average = universal_checkpoint_info.get(
          PARAMETER_TO_AVERAGE_PATTERNS, []
      )
      parameters_with_row_parallelism = universal_checkpoint_info.get(
          PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, []
      )
      vocabulary_parameters = universal_checkpoint_info.get(
          VOCABULARY_PARAMETER_PATTERNS, []
      )
      parameters_with_2_sub_params_cat_dim_0 = (
          universal_checkpoint_info.get(
              PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0, []
          )
      )
      unmatched_patterns = set(
          replicated_parameters
          + parameters_to_average
          + parameters_with_row_parallelism
          + vocabulary_parameters
          + parameters_with_2_sub_params_cat_dim_0
      )
  
      def get_matched_pattern(patterns_, name_):
          matched_ = [
              pattern_
              for pattern_ in patterns_
              if re.match(pattern_, name_)
          ]
          assert (
              len(matched_) <= 1
          ), f"Got more than one matching patterns={matched_} for {name_}"
          if matched_:
              pattern_ = matched_[0]
              unmatched_patterns.discard(pattern_)
              return pattern_
          return None
  
      for state in ("fp32", "exp_avg", "exp_avg_sq"):
          slices = _merge_zero_shards(
              slice_base_path, state, tp_degree, shape
          )
          final_path = os.path.join(param_base_path, f"{state}.pt")
  
          # print(f"Expected shape: {shape}")
          # print(f"Fragment sizes:", list(frag.shape for frag in slices))
          ckpt_dict = {}
          if get_matched_pattern(replicated_parameters, name):
              if len(slices) > 1:
                  assert all(
                      [
                          slices[0].equal(other_slice)
                          for other_slice in slices[1:]
                      ]
                  )
              param = slices[0]
              # print(f'replicate {name} using first slice')
          elif get_matched_pattern(parameters_to_average, name):
              param = sum(slices) / len(slices)
              # print(f'merge {name} using average')
          elif get_matched_pattern(
              parameters_with_2_sub_params_cat_dim_0, name
          ):
              cat_dim = 0
              chunked_slices = [
                  torch.chunk(s, 2, dim=cat_dim) for s in slices
              ]
              merged_chunks_0 = torch.cat(
                  [s[0] for s in chunked_slices], dim=cat_dim
              )
              merged_chunks_1 = torch.cat(
                  [s[1] for s in chunked_slices], dim=cat_dim
              )
              param = torch.cat(
                  [merged_chunks_0, merged_chunks_1], dim=cat_dim
              )
              ckpt_dict[CAT_DIM] = cat_dim
              ckpt_dict[PARAM_N_SUB_PARAMS] = 2
          else:
              cat_dim = (
                  1
                  if get_matched_pattern(
                      parameters_with_row_parallelism, name
                  )
                  else 0
              )
              # print(f"merge {name} with CAT DIM: {cat_dim}")
              param = torch.cat(slices, dim=cat_dim)
              ckpt_dict[CAT_DIM] = cat_dim
  
          if get_matched_pattern(vocabulary_parameters, name):
              # print(f"Before {param.shape=}")
              # strip padding
              original_vocab_size = universal_checkpoint_info[
                  "original_vocab_size"
              ]
              param = param[:original_vocab_size, :]
              ckpt_dict[VOCAB_TENSOR] = True
              # print(f"After {param.shape=}")
  
          # print(f"Final shape: {param.shape}")
          ckpt_dict[PARAM] = param
          _save_checkpoint(final_path, ckpt_dict)
  
      return unmatched_patterns
  
imports: from functools import partial
  import argparse
  import glob
  import itertools
  import multiprocessing
  import os
  import re
  import shutil
  import torch
  import tqdm
  from deepspeed.checkpoint import DeepSpeedCheckpoint
  from deepspeed.checkpoint import OPTIMIZER_STATE_DICT, BASE_OPTIMIZER_STATE, SINGLE_PARTITION_OF_FP32_GROUPS, PARAM_SLICE_MAPPINGS, PARAM_SHAPES, PARAM, CAT_DIM, PARAM_N_SUB_PARAMS, VOCAB_TENSOR, UNIVERSAL_CHECKPOINT_INFO, VOCABULARY_PARAMETER_PATTERNS, PIPELINE_REPLICATED_PARAMETER_PATTERNS, TP_REPLICATED_PARAMETER_PATTERNS, PARAMETER_TO_AVERAGE_PATTERNS, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0
label_desc: Use image node_id 7 for calling a global function with example usage: merge_tp_slices(ds_checkpoint, dir, slice_dir, tp_degree, name_and_shape) and returns: unmatched_patterns, None, pattern_