Unverified Commit f9c0217d authored by Quan (Andy) Gan's avatar Quan (Andy) Gan Committed by GitHub
Browse files

[Bugfix] Fix number of nodes mismatch in to_homo() conversion (#874)

parent bf8bb58f
...@@ -430,9 +430,11 @@ def to_homo(G): ...@@ -430,9 +430,11 @@ def to_homo(G):
eids = [] eids = []
ntype_ids = [] ntype_ids = []
nids = [] nids = []
total_num_nodes = 0
for ntype_id, ntype in enumerate(G.ntypes): for ntype_id, ntype in enumerate(G.ntypes):
num_nodes = G.number_of_nodes(ntype) num_nodes = G.number_of_nodes(ntype)
total_num_nodes += num_nodes
ntype_ids.append(F.full_1d(num_nodes, ntype_id, F.int64, F.cpu())) ntype_ids.append(F.full_1d(num_nodes, ntype_id, F.int64, F.cpu()))
nids.append(F.arange(0, num_nodes)) nids.append(F.arange(0, num_nodes))
...@@ -445,7 +447,7 @@ def to_homo(G): ...@@ -445,7 +447,7 @@ def to_homo(G):
etype_ids.append(F.full_1d(num_edges, etype_id, F.int64, F.cpu())) etype_ids.append(F.full_1d(num_edges, etype_id, F.int64, F.cpu()))
eids.append(F.arange(0, num_edges)) eids.append(F.arange(0, num_edges))
retg = graph((F.cat(srcs, 0), F.cat(dsts, 0))) retg = graph((F.cat(srcs, 0), F.cat(dsts, 0)), card=total_num_nodes)
retg.ndata[NTYPE] = F.cat(ntype_ids, 0) retg.ndata[NTYPE] = F.cat(ntype_ids, 0)
retg.ndata[NID] = F.cat(nids, 0) retg.ndata[NID] = F.cat(nids, 0)
retg.edata[ETYPE] = F.cat(etype_ids, 0) retg.edata[ETYPE] = F.cat(etype_ids, 0)
......
...@@ -652,6 +652,11 @@ def test_convert(): ...@@ -652,6 +652,11 @@ def test_convert():
assert hg.number_of_edges(('user', 'watches', 'movie')) == 1 assert hg.number_of_edges(('user', 'watches', 'movie')) == 1
assert len(hg.etypes) == 2 assert len(hg.etypes) == 2
# hetero_to_homo test case 2
hg = dgl.bipartite([(0, 0), (1, 1)], card=(2, 3))
g = dgl.to_homo(hg)
assert g.number_of_nodes() == 5
def test_subgraph(): def test_subgraph():
g = create_test_heterograph() g = create_test_heterograph()
x = F.randn((3, 5)) x = F.randn((3, 5))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment