def replace_dp(
    graph: DataPipeGraph, old_datapipe: DataPipe, new_datapipe: DataPipe
) -> DataPipeGraph:
    r"""
    Given the graph of DataPipe generated by ``traverse_dps`` function and the
    DataPipe to be replaced and the new DataPipe, return the new graph of
    DataPipe.
    """
    assert len(graph) == 1

    if id(old_datapipe) in graph:
        graph = traverse_dps(new_datapipe)

    final_datapipe = list(graph.values())[0][0]

    for recv_dp, send_graph in graph.values():
        _replace_dp(recv_dp, send_graph, old_datapipe, new_datapipe)

    return traverse_dps(final_datapipe)

    Given the graph of DataPipe generated by ``traverse_dps`` function and the
    DataPipe to be replaced and the new DataPipe, return the new graph of
    DataPipe.
    
python/dgl/graphbolt/datapipes/utils.py