code: def shape_consistency_pass(gm: torch.fx.GraphModule):
      mod_graph = gm.graph
      nodes = tuple(mod_graph.nodes)
      input_dict_node = None
      origin_dict_node = None
  
      # mapping the node into the origin graph index
      node_to_index_dict = {}
      index = 0
      for node in nodes:
          if node.target == "sharding_spec_convert_dict":
              input_dict_node = node
              continue
          if node.target == "origin_node_sharding_spec_dict":
              origin_dict_node = node
              continue
          if not hasattr(node, "best_strategy"):
              continue
          node_to_index_dict[node] = index
          index += 1
      assert input_dict_node is not None
  
      # add shape consistency apply function into graph
      for node in nodes:
          if not hasattr(node, "best_strategy"):
              continue
          with mod_graph.inserting_after(node):
              origin_spec_node = mod_graph.create_node(
                  "call_function",
                  operator.getitem,
                  args=(origin_dict_node, node_to_index_dict[node]),
              )
          with mod_graph.inserting_after(origin_spec_node):
              set_sharding_spec_node = mod_graph.create_node(
                  "call_function",
                  builtins.setattr,
                  args=(node, "sharding_spec", origin_spec_node),
              )
  
          for user_node in node.strategies_vector.successor_nodes:
              node_index = (
                  user_node.strategies_vector.predecessor_nodes.index(
                      node
                  )
              )
              with mod_graph.inserting_before(user_node):
                  input_specs_node = mod_graph.create_node(
                      "call_function",
                      operator.getitem,
                      args=(input_dict_node, node_to_index_dict[node]),
                  )
              with mod_graph.inserting_before(user_node):
                  sharding_spec_node = mod_graph.create_node(
                      "call_function",
                      operator.getitem,
                      args=(input_specs_node, node_index),
                  )
              with mod_graph.inserting_before(user_node):
                  shape_consistency_node = mod_graph.create_node(
                      "call_function",
                      apply,
                      args=(node, sharding_spec_node),
                  )
  
      return gm
  
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 3 for calling a global function with example usage: shape_consistency_pass(gm) and returns: gm