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)
docstr: This method will generate edge_cost for adjacent node pair. Additionally, 'parents' and 'children' attribute will be
  set to node.
code: def _build_cost_graph(self):
      """
      This method will generate edge_cost for adjacent node pair. Additionally, 'parents' and 'children' attribute will be
      set to node.
      """
      self.edge_costs = {}
      if self.simplify:
          self.merge_pair = []
      for strategies_vector in self.leaf_strategies:
          # build edge_cost
          dst_node = strategies_vector.node
          for src_node in strategies_vector.predecessor_nodes:
              if src_node not in self.nodes:
                  continue
              node_pair = (src_node, dst_node)
              edge_cost = {}
              for i in range(len(strategies_vector)):
                  for j in range(len(src_node.strategies_vector)):
                      resharding_cost_item = strategies_vector[
                          i
                      ].resharding_costs[src_node][j]
                      if self.forward_only:
                          edge_cost[(j, i)] = resharding_cost_item.fwd
                      else:
                          edge_cost[(j, i)] = resharding_cost_item.total
              self.edge_costs[node_pair] = edge_cost
          parent_nodes = []
          children_nodes = []
  
          def _check_tensor_in_node(data):
              """
              This method is used to check whether the data has a tensor inside or not.
              """
              has_tensor_flag = False
              if isinstance(data, torch.Tensor):
                  return True
              elif isinstance(data, (tuple, list)):
                  for d in data:
                      has_tensor_flag = (
                          has_tensor_flag or _check_tensor_in_node(d)
                      )
              return has_tensor_flag
  
          for node in strategies_vector.predecessor_nodes:
              if _check_tensor_in_node(node._meta_data):
                  parent_nodes.append(node)
          for node in strategies_vector.successor_nodes:
              if _check_tensor_in_node(node._meta_data):
                  children_nodes.append(node)
  
          setattr(dst_node, "parents", parent_nodes)
          setattr(dst_node, "children", children_nodes)
  
          if self.simplify and strategies_vector.check_merge():
              for followed_node in strategies_vector.predecessor_nodes:
                  # we only merge node pairs which src node has a tensor element inside.
                  # This is necessary because the node without a tensor element inside will not
                  # be assigned any strategy.
                  if _check_tensor_in_node(followed_node._meta_data):
                      self.merge_pair.append((followed_node, dst_node))
  
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 3 for calling the CostGraph obj's underlying member method code with example usage: obj._build_cost_graph() and returns: has_tensor_flag, True