code: def do_dp_split_gpipe(
      node_list, compute_costs, num_stages: int, num_microbatches: int
  ):
      # Ignore the memory cost profiling in Alpa's design for convenience.
      max_compute_costs = np.sort(np.unique(compute_costs))
      best_cost = np.inf
      best_solution = None
      last_max_compute_cost = 0.0
      gap = 1e6  # temporary magic number, unit: flops
  
      for max_compute_cost in tqdm.tqdm(max_compute_costs):
          # Pruning to reduce search space.
          if max_compute_cost * num_microbatches >= best_cost:
              break
          if max_compute_cost - last_max_compute_cost < gap:
              continue
  
          cost, solution = do_dp_split_gpipe_impl(
              len(node_list),
              num_stages,
              num_microbatches,
              compute_costs,
              max_compute_cost,
          )
  
          if cost < best_cost:
              best_cost = cost
              best_solution = solution
          last_max_compute_cost = max_compute_cost
      return best_cost, best_solution
  
file_path: ColossalAI/colossalai/fx/passes/adding_split_node_pass.py
imports: import numpy
  import torch
  import tqdm
  from colossalai.fx.passes.split_module import split_module
label_desc: Use image node_id 7 for calling a global function with example usage: do_dp_split_gpipe(node_list, compute_costs, num_stages, num_microbatches) and returns: best_cost,  best_solution