mirror of
https://github.com/dmlc/dgl.git
synced 2026-06-04 19:44:23 +08:00
8a83027274
* merge * format * rename * sort * sort * update * update * update * Update tests/utils/checks.py Co-authored-by: Mufei Li <mufeili1996@gmail.com> --------- Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal> Co-authored-by: Mufei Li <mufeili1996@gmail.com>
56 lines
2.2 KiB
Python
56 lines
2.2 KiB
Python
import unittest
|
|
|
|
import backend as F
|
|
|
|
from dgl.distributed import graph_partition_book as gpb
|
|
from dgl.partition import NDArrayPartition
|
|
from utils import parametrize_idtype
|
|
|
|
|
|
@unittest.skipIf(
|
|
F._default_context_str == "cpu",
|
|
reason="NDArrayPartition only works on GPU.",
|
|
)
|
|
@parametrize_idtype
|
|
def test_get_node_partition_from_book(idtype):
|
|
node_map = {"_N": F.tensor([[0, 3], [4, 5], [6, 10]], dtype=idtype)}
|
|
edge_map = {
|
|
("_N", "_E", "_N"): F.tensor([[0, 9], [10, 15], [16, 25]], dtype=idtype)
|
|
}
|
|
ntypes = {ntype: i for i, ntype in enumerate(node_map)}
|
|
etypes = {etype: i for i, etype in enumerate(edge_map)}
|
|
book = gpb.RangePartitionBook(0, 3, node_map, edge_map, ntypes, etypes)
|
|
partition = gpb.get_node_partition_from_book(book, F.ctx())
|
|
assert partition.num_parts() == 3
|
|
assert partition.array_size() == 11
|
|
|
|
# Test map_to_local
|
|
test_ids = F.copy_to(F.tensor([0, 2, 6, 7, 10], dtype=idtype), F.ctx())
|
|
act_ids = partition.map_to_local(test_ids)
|
|
exp_ids = F.copy_to(F.tensor([0, 2, 0, 1, 4], dtype=idtype), F.ctx())
|
|
assert F.array_equal(act_ids, exp_ids)
|
|
|
|
# Test map_to_global
|
|
test_ids = F.copy_to(F.tensor([0, 2], dtype=idtype), F.ctx())
|
|
act_ids = partition.map_to_global(test_ids, 0)
|
|
exp_ids = F.copy_to(F.tensor([0, 2], dtype=idtype), F.ctx())
|
|
assert F.array_equal(act_ids, exp_ids)
|
|
|
|
test_ids = F.copy_to(F.tensor([0, 1], dtype=idtype), F.ctx())
|
|
act_ids = partition.map_to_global(test_ids, 1)
|
|
exp_ids = F.copy_to(F.tensor([4, 5], dtype=idtype), F.ctx())
|
|
assert F.array_equal(act_ids, exp_ids)
|
|
|
|
test_ids = F.copy_to(F.tensor([0, 1, 4], dtype=idtype), F.ctx())
|
|
act_ids = partition.map_to_global(test_ids, 2)
|
|
exp_ids = F.copy_to(F.tensor([6, 7, 10], dtype=idtype), F.ctx())
|
|
assert F.array_equal(act_ids, exp_ids)
|
|
|
|
# Test generate_permutation
|
|
test_ids = F.copy_to(F.tensor([6, 0, 7, 2, 10], dtype=idtype), F.ctx())
|
|
perm, split_sum = partition.generate_permutation(test_ids)
|
|
exp_perm = F.copy_to(F.tensor([1, 3, 0, 2, 4], dtype=idtype), F.ctx())
|
|
exp_sum = F.copy_to(F.tensor([2, 0, 3]), F.ctx())
|
|
assert F.array_equal(perm, exp_perm)
|
|
assert F.array_equal(split_sum, exp_sum)
|