def METHOD_NAME():
    chunks = [
        TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None).data for n in range(3)
    ]
    graph = ChunkGraph([chunks[2]])
    for c in chunks:
        graph.add_node(c)
    chunks[1].op._inputs = [chunks[0]]
    graph.add_edge(chunks[0], chunks[1])
    chunks[2].op._inputs = [chunks[0]]
    graph.add_edge(chunks[0], chunks[2])
    all_bands = [("127.0.0.1", "0"), ("127.0.0.1", "1")]
    chunk_to_bands = {
        chunks[0]: all_bands[0],
    }
    coloring = Coloring(graph, all_bands, chunk_to_bands)
    chunk_to_colors = coloring.color()
    assert len(set(chunk_to_colors.values())) == 1
    assert (
        chunk_to_colors[chunks[0]]
        == chunk_to_colors[chunks[1]]
        == chunk_to_colors[chunks[2]]
    )
    coloring = Coloring(
        graph, all_bands, chunk_to_bands, as_broadcaster_successor_num=1
    )
    chunk_to_colors = coloring.color()
    assert len(set(chunk_to_colors.values())) == 3
    assert (
        len(
            {
                chunk_to_colors[chunks[0]],
                chunk_to_colors[chunks[1]],
                chunk_to_colors[chunks[2]],
            }
        )
        == 3
    )