code: def test_simple_coloring():
      # graph: https://user-images.githubusercontent.com/357506/132340029-b595afcf-3cec-44cb-b1c3-aac379e2e607.png
      chunks = [
          TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None).data
          for n in range(8)
      ]
      graph = ChunkGraph([chunks[3], chunks[7]])
      for c in chunks:
          graph.add_node(c)
      chunks[2].op._inputs = [chunks[0], chunks[1]]
      graph.add_edge(chunks[0], chunks[2])
      graph.add_edge(chunks[1], chunks[2])
      chunks[3].op._inputs = [chunks[2]]
      graph.add_edge(chunks[2], chunks[3])
      chunks[6].op._inputs = [chunks[4], chunks[5]]
      graph.add_edge(chunks[4], chunks[6])
      graph.add_edge(chunks[5], chunks[6])
      chunks[7].op._inputs = [chunks[6]]
      graph.add_edge(chunks[6], chunks[7])
  
      all_bands = [("127.0.0.1", "0"), ("127.0.0.1", "1")]
      chunk_to_bands = {
          chunks[0]: all_bands[0],
          chunks[1]: all_bands[0],
          chunks[4]: all_bands[1],
          chunks[5]: all_bands[1],
      }
  
      # allocate node 0, 1 with band 0, node 4, 5 with band 1
      coloring = Coloring(graph, all_bands, chunk_to_bands)
      chunk_to_colors = coloring.color()
      assert len(set(chunk_to_colors.values())) == 2
      assert (
          chunk_to_colors[chunks[0]]
          == chunk_to_colors[chunks[1]]
          == chunk_to_colors[chunks[2]]
          == chunk_to_colors[chunks[3]]
      )
      assert (
          chunk_to_colors[chunks[4]]
          == chunk_to_colors[chunks[5]]
          == chunk_to_colors[chunks[6]]
          == chunk_to_colors[chunks[7]]
      )
  
      # initial nodes all have different colors
      coloring = Coloring(
          graph, all_bands, chunk_to_bands, initial_same_color_num=1
      )
      chunk_to_colors = coloring.color()
      assert len(set(chunk_to_colors.values())) == 6
      assert (
          len(
              {
                  chunk_to_colors[chunks[0]],
                  chunk_to_colors[chunks[1]],
                  chunk_to_colors[chunks[2]],
              }
          )
          == 3
      )
      assert chunk_to_colors[chunks[2]] == chunk_to_colors[chunks[3]]
  
imports: from .....core import ChunkGraph
  from .....tensor.arithmetic import TensorTreeAdd
  from ..fusion import Coloring
label_desc: Use image node_id 1 for calling a global function with example usage: test_simple_coloring() without return types