def METHOD_NAME():
    # 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]]