code: def _merge_zero_shards(
      param_base_path, state, tp_degree, slice_shape
  ):
      slices = []
      for tp_index in range(tp_degree):
          prefix_path = os.path.join(
              param_base_path, str(tp_index), f"{state}"
          )
          paths = sorted(list(glob.glob(f"{prefix_path}.*")))
          shards = [torch.load(p) for p in paths]
          slice = torch.cat(shards, dim=0).reshape(slice_shape)
          slices.append(slice)
      return slices
  
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 6 for calling a global function with example usage: _merge_zero_shards(param_base_path, state, tp_degree, slice_shape) and returns: slices