class_docstr: Coloring a chunk graph according to an algorithm
  described in https://github.com/mars-project/mars/issues/2435
class_docstr_tok: ["Coloring","a","chunk","graph","according","to","an","algorithm","described","in","https",":","\/\/github.com\/mars-project\/mars\/issues\/2435"]
code: def color(self) -> Dict[ChunkType, int]:
      chunk_to_colors = dict()
  
      # step 1: Coloring the initial nodes according to the bands that assigned by assigner
      op_to_colors = self._color_init_nodes()
  
      # step2: Propagate color in the topological order,
      # if the input nodes have same color, color it with the same color;
      # otherwise, color with a new color.
      broadcaster_chunk_set = set()
      for chunk in self.chunk_graph.topological_iter():
          if (
              self.chunk_graph.count_successors(chunk)
              > self.successor_same_color_num
          ):
              # is broadcaster
              broadcaster_chunk_set.add(chunk)
  
          if chunk.op in op_to_colors:
              # colored
              chunk_to_colors[chunk] = op_to_colors[chunk.op]
              continue
  
          predecessors = self.chunk_graph.predecessors(chunk)
          pred_colors = {op_to_colors[pred.op] for pred in predecessors}
          if (
              len(predecessors) == 1
              and predecessors[0] in broadcaster_chunk_set
          ):
              # TODO: handle situation that chunks which specify reassign_workers
              # predecessor is broadcaster, just allocate a new color
              color = self.next_color()
          elif len(pred_colors) == 1:
              if self._can_color_same(chunk, predecessors):
                  # predecessors have only 1 color, will color with same one
                  color = next(iter(pred_colors))
              else:
                  color = self.next_color()
          else:
              # has more than 1 color, color a new one
              assert len(pred_colors) > 1
              color = self.next_color()
  
          op_to_colors[chunk.op] = chunk_to_colors[chunk] = color
  
      # step 3: Propagate with reversed topological order,
      # check a node with its inputs, if all inputs have different color with itself, skip;
      # otherwise, if some of inputs have the same color, but some others have different color,
      # color the input nodes with same one with a new color, and propagate to its inputs and so on.
      for chunk in self.chunk_graph.topological_iter(reverse=True):
          pred_colors = {
              op_to_colors[pred.op]
              for pred in self.chunk_graph.iter_successors(chunk)
          }
          chunk_color = chunk_to_colors[chunk]
          if chunk_color in pred_colors and len(pred_colors) > 1:
              # conflict
              # color the successors with new colors
              stack = []
              for succ in self.chunk_graph.iter_successors(chunk):
                  if chunk_to_colors[succ] == chunk_color:
                      new_color = op_to_colors[
                          succ.op
                      ] = self.next_color()
                      for c in succ.op.outputs:
                          if (
                              c not in self.chunk_graph
                          ):  # pragma: no cover
                              continue
                          chunk_to_colors[c] = new_color
                          stack.extend(self.chunk_graph.successors(c))
              # color the descendants with same color to the new one
              # the descendants will not be visited more than 2 times
              while len(stack) > 0:
                  node = stack.pop()
                  node_color = chunk_to_colors[node]
                  if node_color == chunk_color:
                      # same color, recolor to the new one
                      node_pred_colors = list(
                          {
                              op_to_colors[inp.op]
                              for inp in self.chunk_graph.iter_predecessors(
                                  node
                              )
                          }
                      )
                      node_input_same_color = len(node_pred_colors) == 1
                      if node_input_same_color:
                          node_new_color = node_pred_colors[0]
                      else:
                          node_new_color = self.next_color()
                      op_to_colors[node.op] = node_new_color
                      for c in node.op.outputs:
                          if (
                              c not in self.chunk_graph
                          ):  # pragma: no cover
                              continue
                          chunk_to_colors[c] = node_new_color
                          stack.extend(self.chunk_graph.successors(c))
  
      return chunk_to_colors
  
file_path: mars/mars/services/task/analyzer/fusion.py
imports: import itertools
  from collections import defaultdict
  from typing import Dict, List
  from ....config import options
  from ....core import ChunkGraph
  from ....core.operand import VirtualOperand
  from ....typing import BandType, ChunkType, OperandType
label_desc: Use image node_id 5 for calling the Coloring obj's underlying member method code with example usage: obj.color() and returns: chunk_to_colors