code: def solution_annotation_pass(
      gm: torch.fx.GraphModule, solution: List[int], device_mesh
  ):
      mod_graph = gm.graph
      nodes = tuple(mod_graph.nodes)
  
      # the dict to get origin sharding spec of node
      origin_node_sharding_spec_dict = {}
      for node_index, (node, strategy_index) in enumerate(
          zip(nodes, solution)
      ):
          strategies_vector = node.strategies_vector
          setattr(
              node, "best_strategy", strategies_vector[strategy_index]
          )
          setattr(
              node,
              "sharding_spec",
              strategies_vector[strategy_index].output_sharding_spec,
          )
          origin_node_sharding_spec_dict[
              node_index
          ] = strategies_vector[strategy_index].output_sharding_spec
  
      # apply the sharding spec of parameters
      for node in nodes:
          if node.op == "call_module":
              target_module = node.graph.owning_module.get_submodule(
                  node.target
              )
              origin_sharding_spec = ShardingSpec(
                  device_mesh, target_module.weight.shape, {}
              )
              setattr(
                  target_module.weight,
                  "sharding_spec",
                  origin_sharding_spec,
              )
              target_weight_sharding_spec = (
                  node.best_strategy.input_shardings[1]
              )
              target_module.weight.data = (
                  target_module.weight.data.permute((1, 0, 2, 3))
              )
              apply(target_module.weight, target_weight_sharding_spec)
              target_module.weight.data = (
                  target_module.weight.data.permute((1, 0, 2, 3))
              )
  
      # the dict to get input sharding specs of user node
      sharding_spec_convert_dict = {}
      for index, node in enumerate(nodes):
          target_sharding_specs = []
          for user_node in node.strategies_vector.successor_nodes:
              node_index = (
                  user_node.strategies_vector.predecessor_nodes.index(
                      node
                  )
              )
              target_sharding_spec = (
                  user_node.best_strategy.input_shardings[node_index]
              )
              target_sharding_specs.append(target_sharding_spec)
          sharding_spec_convert_dict[index] = target_sharding_specs
  
      # add above dicts into graph
      for node in nodes:
          if node.op != "placeholder":
              with mod_graph.inserting_before(node):
                  input_specs_node = mod_graph.create_node(
                      "placeholder", target="sharding_spec_convert_dict"
                  )
                  origin_specs_node = mod_graph.create_node(
                      "placeholder",
                      target="origin_node_sharding_spec_dict",
                  )
              break
  
      return sharding_spec_convert_dict, origin_node_sharding_spec_dict
  
file_path: ColossalAI/colossalai/fx/passes/experimental/adding_shape_consistency_pass.py
imports: import builtins
  import operator
  from typing import List
  import torch
  from colossalai.tensor.shape_consistency import ShapeConsistencyManager
  from colossalai.tensor.sharding_spec import ShardingSpec
label_desc: Use image node_id 2 for calling a global function with example usage: solution_annotation_pass(gm, solution, device_mesh) and returns: sharding_spec_convert_dict,  origin_node_sharding_spec_dict