class_docstr: A graph data structure to simplify the edge cost graph. It has two main functions:
  1. To feed the quadratic resharding costs into solver, we need to linearize it. We build edge_cost in
  CostGraph, and it stored every combinations of strategies for a src-dst node pair in an 1D list.
  2. To reduce the searching space, we merge computationally-trivial operators, such as
  element-wise operators, transpose, and reduction, into their following nodes. The merging information will
  be given by the StrategiesVector depending on the type of target node and following nodes.
  
  Argument:
      leaf_strategies(List[StrategiesVector]): It stores StrategiesVector of every nodes on the graph.
      simplify(bool, optional): The generated cost graph will be simplified if it is true. (default to True)
code: def __init__(
      self, leaf_strategies, simplify=True, forward_only=False
  ):
      self.leaf_strategies = leaf_strategies
      self.nodes = [
          strategies_vector.node
          for strategies_vector in self.leaf_strategies
      ]
      # stores number of strategies in each node
      self.node_lens = {
          strategies_vector.node: len(strategies_vector)
          for strategies_vector in self.leaf_strategies
      }
      # extra_node_costs will store the extra costs introduced by merging nodes
      self.extra_node_costs = {}
      self.following_dict = {}
      self.simplify = simplify
      self.forward_only = forward_only
      self._build_cost_graph()
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/solver/cost_graph.py
imports: import torch
  from colossalai.auto_parallel.tensor_shard.constants import INFINITY_COST
label_desc: Use image node_id 1 to create a new CostGraph object  with example: obj = CostGraph(leaf_strategies, simplify, forward_only)