[Refactor] Explicit dtype for HeteroGraph (#1467)

* 111

* 111

* lint

* lint

* lint

* lint

* fix

* lint

* try

* fix

* lint

* lint

* test

* fix

* ttt

* test

* fix

* fix

* fix

* mxnet

* 111

* fix 64bits computation

* pylint

* roll back

* fix

* lint

* fix hetero_from_relations

* remove index_dtype in to_homo and to_hetero

* fix

* fix

* fix

* fix

* remove default

* fix

* lint

* fix

* fix error message

* fix error

* lint

* macro dispatch

* try

* lint

* remove nbits

* error message

* fix

* fix

* lint

* lint

* lint

* fix

* lint

* fix

* fix random walk

* lint

* lint

* fix

* fix

* fix

* lint

* fix

* lint
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>

tests/compute/test_randomwalk.py

@@ -2,6 +2,8 @@ import dgl
 from dgl import utils
 import backend as F
 import numpy as np
+from utils import parametrize_dtype
+import pytest
 
 def test_random_walk():
     edge_list = [(0, 1), (1, 2), (2, 3), (3, 4),
@@ -58,9 +60,10 @@ def test_random_walk_with_restart():
             trace_diff = np.diff(F.zerocopy_to_numpy(t), axis=-1)
             assert (trace_diff % 2 == 0).all()
 
-def test_metapath_random_walk():
-    g1 = dgl.bipartite(([0, 1, 2, 3], [0, 1, 2, 3]), 'a', 'ab', 'b')
-    g2 = dgl.bipartite(([0, 0, 1, 1, 2, 2, 3, 3], [1, 3, 2, 0, 3, 1, 0, 2]), 'b', 'ba', 'a')
+@parametrize_dtype
+def test_metapath_random_walk(index_dtype):
+    g1 = dgl.bipartite(([0, 1, 2, 3], [0, 1, 2, 3]), 'a', 'ab', 'b', index_dtype=index_dtype)
+    g2 = dgl.bipartite(([0, 0, 1, 1, 2, 2, 3, 3], [1, 3, 2, 0, 3, 1, 0, 2]), 'b', 'ba', 'a', index_dtype=index_dtype)
     G = dgl.hetero_from_relations([g1, g2])
     seeds = [0, 1]
     traces = dgl.contrib.sampling.metapath_random_walk(G, ['ab', 'ba'] * 4, seeds, 3)
@@ -74,5 +77,5 @@ def test_metapath_random_walk():
                 assert g2.has_edge_between(trace[2 * i + 1], trace[2 * i + 2])
 
 if __name__ == '__main__':
-    test_random_walk()
-    test_metapath_random_walk()
+    # test_random_walk()
+    test_metapath_random_walk("int32")

tests/compute/test_transform.py

@@ -7,6 +7,7 @@ import dgl.function as fn
 import backend as F
 from dgl.graph_index import from_scipy_sparse_matrix
 import unittest
+from utils import parametrize_dtype
 
 D = 5
 
@@ -271,13 +272,15 @@ def test_metis_partition():
         print(g.number_of_edges() - num_inner_edges)
 
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
-def test_in_subgraph():
-    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow')
-    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game')
-    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user')
-    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin')
+@parametrize_dtype
+def test_in_subgraph(index_dtype):
+    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', index_dtype=index_dtype)
+    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', index_dtype=index_dtype)
+    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', index_dtype=index_dtype)
+    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', index_dtype=index_dtype)
     hg = dgl.hetero_from_relations([g1, g2, g3, g4])
     subg = dgl.in_subgraph(hg, {'user' : [0,1], 'game' : 0})
+    assert subg._idtype_str == index_dtype
     assert len(subg.ntypes) == 3
     assert len(subg.etypes) == 4
     u, v = subg['follow'].edges()
@@ -295,13 +298,15 @@ def test_in_subgraph():
     assert subg['flips'].number_of_edges() == 0
 
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
-def test_out_subgraph():
-    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow')
-    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game')
-    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user')
-    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin')
+@parametrize_dtype
+def test_out_subgraph(index_dtype):
+    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', index_dtype=index_dtype)
+    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', index_dtype=index_dtype)
+    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', index_dtype=index_dtype)
+    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', index_dtype=index_dtype)
     hg = dgl.hetero_from_relations([g1, g2, g3, g4])
     subg = dgl.out_subgraph(hg, {'user' : [0,1], 'game' : 0})
+    assert subg._idtype_str == index_dtype
     assert len(subg.ntypes) == 3
     assert len(subg.etypes) == 4
     u, v = subg['follow'].edges()
@@ -322,20 +327,21 @@ def test_out_subgraph():
     assert F.array_equal(hg['flips'].edge_ids(u, v), subg['flips'].edata[dgl.EID])
 
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented")
-def test_compact():
+@parametrize_dtype
+def test_compact(index_dtype):
     g1 = dgl.heterograph({
         ('user', 'follow', 'user'): [(1, 3), (3, 5)],
         ('user', 'plays', 'game'): [(2, 4), (3, 4), (2, 5)],
         ('game', 'wished-by', 'user'): [(6, 7), (5, 7)]},
-        {'user': 20, 'game': 10})
+        {'user': 20, 'game': 10}, index_dtype=index_dtype)
 
     g2 = dgl.heterograph({
         ('game', 'clicked-by', 'user'): [(3, 1)],
         ('user', 'likes', 'user'): [(1, 8), (8, 9)]},
-        {'user': 20, 'game': 10})
+        {'user': 20, 'game': 10}, index_dtype=index_dtype)
 
-    g3 = dgl.graph([(0, 1), (1, 2)], num_nodes=10, ntype='user')
-    g4 = dgl.graph([(1, 3), (3, 5)], num_nodes=10, ntype='user')
+    g3 = dgl.graph([(0, 1), (1, 2)], num_nodes=10, ntype='user', index_dtype=index_dtype)
+    g4 = dgl.graph([(1, 3), (3, 5)], num_nodes=10, ntype='user', index_dtype=index_dtype)
 
     def _check(g, new_g, induced_nodes):
         assert g.ntypes == new_g.ntypes
@@ -358,13 +364,15 @@ def test_compact():
     new_g1 = dgl.compact_graphs(g1)
     induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
     induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
+    assert new_g1._idtype_str == index_dtype
     assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
     assert set(induced_nodes['game']) == set([4, 5, 6])
     _check(g1, new_g1, induced_nodes)
 
     # Test with always_preserve given a dict
     new_g1 = dgl.compact_graphs(
-        g1, always_preserve={'game': F.tensor([4, 7], dtype=F.int64)})
+        g1, always_preserve={'game': F.tensor([4, 7], dtype=getattr(F, index_dtype))})
+    assert new_g1._idtype_str == index_dtype
     induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
     induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
     assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
@@ -373,9 +381,11 @@ def test_compact():
 
     # Test with always_preserve given a tensor
     new_g3 = dgl.compact_graphs(
-        g3, always_preserve=F.tensor([1, 7], dtype=F.int64))
+        g3, always_preserve=F.tensor([1, 7], dtype=getattr(F, index_dtype)))
     induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes}
     induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
+    
+    assert new_g3._idtype_str == index_dtype
     assert set(induced_nodes['user']) == set([0, 1, 2, 7])
     _check(g3, new_g3, induced_nodes)
 
@@ -383,6 +393,8 @@ def test_compact():
     new_g1, new_g2 = dgl.compact_graphs([g1, g2])
     induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
     induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
+    assert new_g1._idtype_str == index_dtype
+    assert new_g2._idtype_str == index_dtype
     assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
     assert set(induced_nodes['game']) == set([3, 4, 5, 6])
     _check(g1, new_g1, induced_nodes)
@@ -390,9 +402,11 @@ def test_compact():
 
     # Test multiple graphs with always_preserve given a dict
     new_g1, new_g2 = dgl.compact_graphs(
-        [g1, g2], always_preserve={'game': F.tensor([4, 7], dtype=F.int64)})
+        [g1, g2], always_preserve={'game': F.tensor([4, 7], dtype=getattr(F, index_dtype))})
     induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
-    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
+    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}    
+    assert new_g1._idtype_str == index_dtype
+    assert new_g2._idtype_str == index_dtype
     assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
     assert set(induced_nodes['game']) == set([3, 4, 5, 6, 7])
     _check(g1, new_g1, induced_nodes)
@@ -400,18 +414,21 @@ def test_compact():
 
     # Test multiple graphs with always_preserve given a tensor
     new_g3, new_g4 = dgl.compact_graphs(
-        [g3, g4], always_preserve=F.tensor([1, 7], dtype=F.int64))
+        [g3, g4], always_preserve=F.tensor([1, 7], dtype=getattr(F, index_dtype)))
     induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes}
     induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
+    
+    assert new_g3._idtype_str == index_dtype
+    assert new_g4._idtype_str == index_dtype
     assert set(induced_nodes['user']) == set([0, 1, 2, 3, 5, 7])
     _check(g3, new_g3, induced_nodes)
     _check(g4, new_g4, induced_nodes)
 
-
-def test_to_simple():
+@parametrize_dtype
+def test_to_simple(index_dtype):
     g = dgl.heterograph({
         ('user', 'follow', 'user'): [(0, 1), (1, 3), (2, 2), (1, 3), (1, 4), (1, 4)],
-        ('user', 'plays', 'game'): [(3, 5), (2, 3), (1, 4), (1, 4), (3, 5), (2, 3), (2, 3)]})
+        ('user', 'plays', 'game'): [(3, 5), (2, 3), (1, 4), (1, 4), (3, 5), (2, 3), (2, 3)]}, index_dtype=index_dtype)
     sg = dgl.to_simple(g, return_counts='weights', writeback_mapping='new_eid')
 
     for etype in g.canonical_etypes:
@@ -434,7 +451,8 @@ def test_to_simple():
             assert eid_map[i] == suv.index(e)
 
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented")
-def test_to_block():
+@parametrize_dtype
+def test_to_block(index_dtype):
     def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True):
         if dst_nodes is not None:
             assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes)
@@ -471,7 +489,7 @@ def test_to_block():
     g = dgl.heterograph({
         ('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)],
         ('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)],
-        ('B', 'BA', 'A'): [(2, 3), (3, 2)]})
+        ('B', 'BA', 'A'): [(2, 3), (3, 2)]}, index_dtype=index_dtype)
     g_a = g['AA']
 
     bg = dgl.to_block(g_a)
@@ -484,39 +502,41 @@ def test_to_block():
     assert bg.number_of_src_nodes() == 4
     assert bg.number_of_dst_nodes() == 4
 
-    dst_nodes = F.tensor([3, 4], dtype=F.int64)
+    dst_nodes = F.tensor([3, 4], dtype=getattr(F, index_dtype))
     bg = dgl.to_block(g_a, dst_nodes)
     check(g_a, bg, 'A', 'AA', dst_nodes)
 
-    dst_nodes = F.tensor([4, 3, 2, 1], dtype=F.int64)
+    dst_nodes = F.tensor([4, 3, 2, 1], dtype=getattr(F, index_dtype))
     bg = dgl.to_block(g_a, dst_nodes)
     check(g_a, bg, 'A', 'AA', dst_nodes)
 
     g_ab = g['AB']
 
     bg = dgl.to_block(g_ab)
+    assert bg._idtype_str == index_dtype
     assert bg.number_of_nodes('SRC/B') == 4
     assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID])
     assert bg.number_of_nodes('DST/A') == 0
     checkall(g_ab, bg, None)
 
-    dst_nodes = {'B': F.tensor([5, 6], dtype=F.int64)}
+    dst_nodes = {'B': F.tensor([5, 6], dtype=getattr(F, index_dtype))}
     bg = dgl.to_block(g, dst_nodes)
     assert bg.number_of_nodes('SRC/B') == 2
     assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID])
     assert bg.number_of_nodes('DST/A') == 0
     checkall(g, bg, dst_nodes)
 
-    dst_nodes = {'A': F.tensor([3, 4], dtype=F.int64), 'B': F.tensor([5, 6], dtype=F.int64)}
+    dst_nodes = {'A': F.tensor([3, 4], dtype=getattr(F, index_dtype)), 'B': F.tensor([5, 6], dtype=getattr(F, index_dtype))}
     bg = dgl.to_block(g, dst_nodes)
     checkall(g, bg, dst_nodes)
 
-    dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=F.int64), 'B': F.tensor([3, 5, 6, 1], dtype=F.int64)}
+    dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=getattr(F, index_dtype)), 'B': F.tensor([3, 5, 6, 1], dtype=getattr(F, index_dtype))}
     bg = dgl.to_block(g, dst_nodes=dst_nodes)
     checkall(g, bg, dst_nodes)
 
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
-def test_remove_edges():
+@parametrize_dtype
+def test_remove_edges(index_dtype):
     def check(g1, etype, g, edges_removed):
         src, dst, eid = g.edges(etype=etype, form='all')
         src1, dst1 = g1.edges(etype=etype, order='eid')
@@ -535,34 +555,35 @@ def test_remove_edges():
         for s, d, e in zip(src1, dst1, eid1):
             assert (s, d, e) in sde_set
         assert not np.isin(edges_removed, eid1).any()
+        assert g1.idtype == g.idtype
 
     for fmt in ['coo', 'csr', 'csc']:
         for edges_to_remove in [[2], [2, 2], [3, 2], [1, 3, 1, 2]]:
-            g = dgl.graph([(0, 1), (2, 3), (1, 2), (3, 4)], restrict_format=fmt)
-            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove))
+            g = dgl.graph([(0, 1), (2, 3), (1, 2), (3, 4)], restrict_format=fmt, index_dtype=index_dtype)
+            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, getattr(F, index_dtype)))
             check(g1, None, g, edges_to_remove)
 
             g = dgl.graph(
                 spsp.csr_matrix(([1, 1, 1, 1], ([0, 2, 1, 3], [1, 3, 2, 4])), shape=(5, 5)),
-                restrict_format=fmt)
-            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove))
+                restrict_format=fmt, index_dtype=index_dtype)
+            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, getattr(F, index_dtype)))
             check(g1, None, g, edges_to_remove)
 
     g = dgl.heterograph({
         ('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)],
         ('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)],
-        ('B', 'BA', 'A'): [(2, 3), (3, 2)]})
-    g2 = dgl.remove_edges(g, {'AA': F.tensor([2]), 'AB': F.tensor([3]), 'BA': F.tensor([1])})
+        ('B', 'BA', 'A'): [(2, 3), (3, 2)]}, index_dtype=index_dtype)
+    g2 = dgl.remove_edges(g, {'AA': F.tensor([2], getattr(F, index_dtype)), 'AB': F.tensor([3], getattr(F, index_dtype)), 'BA': F.tensor([1], getattr(F, index_dtype))})
     check(g2, 'AA', g, [2])
     check(g2, 'AB', g, [3])
     check(g2, 'BA', g, [1])
 
-    g3 = dgl.remove_edges(g, {'AA': F.tensor([]), 'AB': F.tensor([3]), 'BA': F.tensor([1])})
+    g3 = dgl.remove_edges(g, {'AA': F.tensor([], getattr(F, index_dtype)), 'AB': F.tensor([3], getattr(F, index_dtype)), 'BA': F.tensor([1], getattr(F, index_dtype))})
     check(g3, 'AA', g, [])
     check(g3, 'AB', g, [3])
     check(g3, 'BA', g, [1])
 
-    g4 = dgl.remove_edges(g, {'AB': F.tensor([3, 1, 2, 0])})
+    g4 = dgl.remove_edges(g, {'AB': F.tensor([3, 1, 2, 0], getattr(F, index_dtype))})
     check(g4, 'AA', g, [])
     check(g4, 'AB', g, [3, 1, 2, 0])
     check(g4, 'BA', g, [])
@@ -594,22 +615,22 @@ def test_cast():
     assert F.array_equal(g2dst, gdst)
 
 if __name__ == '__main__':
-    test_line_graph()
-    test_no_backtracking()
-    test_reverse()
-    test_reverse_shared_frames()
-    test_simple_graph()
-    test_bidirected_graph()
-    test_khop_adj()
-    test_khop_graph()
-    test_laplacian_lambda_max()
-    test_remove_self_loop()
-    test_add_self_loop()
-    test_partition_with_halo()
-    test_metis_partition()
-    test_compact()
-    test_to_simple()
-    test_in_subgraph()
-    test_out_subgraph()
-    test_to_block()
-    test_remove_edges()
+    # test_line_graph()
+    # test_no_backtracking()
+    # test_reverse()
+    # test_reverse_shared_frames()
+    # test_simple_graph()
+    # test_bidirected_graph()
+    # test_khop_adj()
+    # test_khop_graph()
+    # test_laplacian_lambda_max()
+    # test_remove_self_loop()
+    # test_add_self_loop()
+    # test_partition_with_halo()
+    # test_metis_partition()
+    # test_compact()
+    # test_to_simple()
+    # test_in_subgraph("int32")
+    # test_out_subgraph()
+    test_to_block("int32")
+    # test_remove_edges()

tests/compute/utils.py

+import pytest
+parametrize_dtype = pytest.mark.parametrize("index_dtype", ['int32', 'int64'])
 
 def check_fail(fn, *args, **kwargs):
     try: