Unverified Commit 82b11b6c authored by peizhou001's avatar peizhou001 Committed by GitHub
Browse files

[Graphbolt] Patch: Change compact to list (#6183)


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-16-19.ap-northeast-1.compute.internal>
parent 8b23ec60
...@@ -45,7 +45,7 @@ def unique_and_compact( ...@@ -45,7 +45,7 @@ def unique_and_compact(
nodes, empty_tensor, empty_tensor nodes, empty_tensor, empty_tensor
) )
compacted = compacted.split(nums) compacted = compacted.split(nums)
return unique, compacted return unique, list(compacted)
if is_heterogeneous: if is_heterogeneous:
unique, compacted = {}, {} unique, compacted = {}, {}
......
...@@ -28,6 +28,7 @@ def test_unique_and_compact_hetero(): ...@@ -28,6 +28,7 @@ def test_unique_and_compact_hetero():
for ntype, nodes in compacted.items(): for ntype, nodes in compacted.items():
expected_nodes = nodes_dict[ntype] expected_nodes = nodes_dict[ntype]
assert isinstance(nodes, list)
for expected_node, node in zip(expected_nodes, nodes): for expected_node, node in zip(expected_nodes, nodes):
node = unique[ntype][node] node = unique[ntype][node]
assert torch.equal(expected_node, node) assert torch.equal(expected_node, node)
...@@ -42,6 +43,7 @@ def test_unique_and_compact_homo(): ...@@ -42,6 +43,7 @@ def test_unique_and_compact_homo():
assert torch.equal(torch.sort(unique)[0], expected_unique_N) assert torch.equal(torch.sort(unique)[0], expected_unique_N)
assert isinstance(compacted, list)
for expected_node, node in zip(nodes_list, compacted): for expected_node, node in zip(nodes_list, compacted):
node = unique[node] node = unique[node]
assert torch.equal(expected_node, node) assert torch.equal(expected_node, node)
......
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