Function test_partition_with_halo

tests/python/common/transforms/test_transform.py:698–716

()

# Disabled since everything will be on heterogeneous graphs
@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
def test_partition_with_halo():
    g = create_large_graph(1000)
    node_part = np.random.choice(4, g.num_nodes())
    subgs, _, _ = dgl.transforms.partition_graph_with_halo(
        g, node_part, 2, reshuffle=True
    )
    for part_id, subg in subgs.items():
        node_ids = np.nonzero(node_part == part_id)[0]
        lnode_ids = np.nonzero(F.asnumpy(subg.ndata["inner_node"]))[0]
        orig_nids = F.asnumpy(subg.ndata["orig_id"])[lnode_ids]
        assert np.all(np.sort(orig_nids) == node_ids)
        assert np.all(
            F.asnumpy(subg.in_degrees(lnode_ids))
            == F.asnumpy(g.in_degrees(orig_nids))
        )
        assert np.all(
            F.asnumpy(subg.out_degrees(lnode_ids))
            == F.asnumpy(g.out_degrees(orig_nids))
        )


@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")