code: def uniform_split_pass(gm: torch.fx.GraphModule, pp_size: int):
      mod_graph = gm.graph
      valid_children_size = 0
      valid_children = []
      for module in mod_graph.owning_module.children():
          valid_children_size += 1
          valid_children.append(module)
  
      if valid_children_size < pp_size:
          # If valid children is not enough to shard, we will use balanced policy instead of uniform policy.
          return balanced_split_pass(gm, pp_size)
      layers_per_partition = valid_children_size // pp_size
      accumulate_layer_amount = 0
      for node in mod_graph.nodes:
          if pp_size <= 1:
              break
          if node.op == "call_module":
              target_module = node.graph.owning_module.get_submodule(
                  node.target
              )
              if target_module in valid_children:
                  accumulate_layer_amount += 1
          if accumulate_layer_amount == layers_per_partition:
              accumulate_layer_amount = 0
              pp_size -= 1
              with mod_graph.inserting_after(node):
                  split_node = mod_graph.create_node(
                      "call_function", pipe_split
                  )
      gm.recompile()
      return gm
  
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 13 for calling a global function with example usage: uniform_split_pass(gm, pp_size) and returns: gm, balanced_split_pass