Skip to content

GitLab

Unverified Commit 14251504 authored by nv-dlasalle's avatar nv-dlasalle Committed by GitHub
Browse files

Fix test naming (#4023) 


Co-authored-by: default avatarMufei Li <mufeili1996@gmail.com>
parent 7c598aac
Hide whitespace changes
Inline Side-by-side
Showing with 151 additions and 161 deletions
tests/compute/test_sort.py
View file @ 14251504
...@@ -8,7 +8,7 @@ import backend as F ...@@ -8,7 +8,7 @@ import backend as F
import networkx as nx import networkx as nx
import unittest, pytest import unittest, pytest
from dgl import DGLError from dgl import DGLError
from utils import parametrize_dtype from test_utils import parametrize_idtype
def create_test_heterograph(num_nodes, num_adj, idtype): def create_test_heterograph(num_nodes, num_adj, idtype):
if isinstance(num_adj, int): if isinstance(num_adj, int):
...@@ -49,7 +49,7 @@ def check_sort(spm, tag_arr=None, tag_pos=None): ...@@ -49,7 +49,7 @@ def check_sort(spm, tag_arr=None, tag_pos=None):
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sorting by tag not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sorting by tag not implemented")
@parametrize_dtype @parametrize_idtype
def test_sort_with_tag(idtype): def test_sort_with_tag(idtype):
num_nodes, num_adj, num_tags = 200, [20, 50], 5 num_nodes, num_adj, num_tags = 200, [20, 50], 5
g = create_test_heterograph(num_nodes, num_adj, idtype=idtype) g = create_test_heterograph(num_nodes, num_adj, idtype=idtype)
...@@ -68,7 +68,7 @@ def test_sort_with_tag(idtype): ...@@ -68,7 +68,7 @@ def test_sort_with_tag(idtype):
assert(not check_sort(old_csc, tag)) assert(not check_sort(old_csc, tag))
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sorting by tag not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sorting by tag not implemented")
@parametrize_dtype @parametrize_idtype
def test_sort_with_tag_bipartite(idtype): def test_sort_with_tag_bipartite(idtype):
num_nodes, num_adj, num_tags = 200, [20, 50], 5 num_nodes, num_adj, num_tags = 200, [20, 50], 5
g = create_test_heterograph(num_nodes, num_adj, idtype=idtype) g = create_test_heterograph(num_nodes, num_adj, idtype=idtype)
......
tests/compute/test_sparse.py
View file @ 14251504
from dgl.ops import gspmm, gsddmm, edge_softmax, segment_reduce from dgl.ops import gspmm, gsddmm, edge_softmax, segment_reduce
from test_utils.graph_cases import get_cases from test_utils.graph_cases import get_cases
from utils import parametrize_dtype from test_utils import parametrize_idtype
import dgl import dgl
import random import random
import pytest, unittest import pytest, unittest
...@@ -99,7 +99,7 @@ edge_softmax_shapes = [ ...@@ -99,7 +99,7 @@ edge_softmax_shapes = [
@pytest.mark.parametrize('shp', spmm_shapes) @pytest.mark.parametrize('shp', spmm_shapes)
@pytest.mark.parametrize('msg', ['add', 'sub', 'mul', 'div', 'copy_lhs', 'copy_rhs']) @pytest.mark.parametrize('msg', ['add', 'sub', 'mul', 'div', 'copy_lhs', 'copy_rhs'])
@pytest.mark.parametrize('reducer', ['sum', 'min', 'max']) @pytest.mark.parametrize('reducer', ['sum', 'min', 'max'])
@parametrize_dtype @parametrize_idtype
def test_spmm(idtype, g, shp, msg, reducer): def test_spmm(idtype, g, shp, msg, reducer):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
print(g) print(g)
...@@ -159,7 +159,7 @@ def test_spmm(idtype, g, shp, msg, reducer): ...@@ -159,7 +159,7 @@ def test_spmm(idtype, g, shp, msg, reducer):
@pytest.mark.parametrize('lhs_target', ['u', 'v', 'e']) @pytest.mark.parametrize('lhs_target', ['u', 'v', 'e'])
@pytest.mark.parametrize('rhs_target', ['u', 'v', 'e']) @pytest.mark.parametrize('rhs_target', ['u', 'v', 'e'])
@pytest.mark.parametrize('msg', ['add', 'sub', 'mul', 'div', 'dot', 'copy_lhs', 'copy_rhs']) @pytest.mark.parametrize('msg', ['add', 'sub', 'mul', 'div', 'dot', 'copy_lhs', 'copy_rhs'])
@parametrize_dtype @parametrize_idtype
def test_sddmm(g, shp, lhs_target, rhs_target, msg, idtype): def test_sddmm(g, shp, lhs_target, rhs_target, msg, idtype):
if lhs_target == rhs_target: if lhs_target == rhs_target:
return return
...@@ -229,7 +229,7 @@ def test_sddmm(g, shp, lhs_target, rhs_target, msg, idtype): ...@@ -229,7 +229,7 @@ def test_sddmm(g, shp, lhs_target, rhs_target, msg, idtype):
@pytest.mark.parametrize('g', get_cases(['clique'])) @pytest.mark.parametrize('g', get_cases(['clique']))
@pytest.mark.parametrize('norm_by', ['src', 'dst']) @pytest.mark.parametrize('norm_by', ['src', 'dst'])
@pytest.mark.parametrize('shp', edge_softmax_shapes) @pytest.mark.parametrize('shp', edge_softmax_shapes)
@parametrize_dtype @parametrize_idtype
def test_edge_softmax(g, norm_by, shp, idtype): def test_edge_softmax(g, norm_by, shp, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
edata = F.tensor(np.random.rand(g.number_of_edges(), *shp)) edata = F.tensor(np.random.rand(g.number_of_edges(), *shp))
...@@ -288,7 +288,7 @@ def test_segment_reduce(reducer): ...@@ -288,7 +288,7 @@ def test_segment_reduce(reducer):
print('backward passed') print('backward passed')
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256]) @pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256])
def test_segment_mm(idtype, feat_size): def test_segment_mm(idtype, feat_size):
import torch import torch
...@@ -323,7 +323,7 @@ def test_segment_mm(idtype, feat_size): ...@@ -323,7 +323,7 @@ def test_segment_mm(idtype, feat_size):
assert torch.allclose(db, db_t, atol=1e-4, rtol=1e-4) assert torch.allclose(db, db_t, atol=1e-4, rtol=1e-4)
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256]) @pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256])
def test_gather_mm_idx_b(idtype, feat_size): def test_gather_mm_idx_b(idtype, feat_size):
import torch import torch
...@@ -353,7 +353,7 @@ def test_gather_mm_idx_b(idtype, feat_size): ...@@ -353,7 +353,7 @@ def test_gather_mm_idx_b(idtype, feat_size):
assert torch.allclose(db, db_t, atol=1e-4, rtol=1e-4) assert torch.allclose(db, db_t, atol=1e-4, rtol=1e-4)
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256]) @pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256])
def _test_gather_mm_idx_a(idtype, feat_size): def _test_gather_mm_idx_a(idtype, feat_size):
# TODO(minjie): currently disabled due to bugs in the CUDA kernel. Need to fix it later. # TODO(minjie): currently disabled due to bugs in the CUDA kernel. Need to fix it later.
......
tests/compute/test_specialization.py
View file @ 14251504
...@@ -3,7 +3,7 @@ import scipy.sparse as sp ...@@ -3,7 +3,7 @@ import scipy.sparse as sp
import dgl import dgl
import dgl.function as fn import dgl.function as fn
import backend as F import backend as F
from test_utils import parametrize_dtype from test_utils import parametrize_idtype
D = 5 D = 5
...@@ -22,7 +22,7 @@ def generate_graph(idtype): ...@@ -22,7 +22,7 @@ def generate_graph(idtype):
g.edata.update({'e1': weights, 'e2': F.unsqueeze(weights, 1)}) g.edata.update({'e1': weights, 'e2': F.unsqueeze(weights, 1)})
return g return g
@parametrize_dtype @parametrize_idtype
def test_v2v_update_all(idtype): def test_v2v_update_all(idtype):
def _test(fld): def _test(fld):
def message_func(edges): def message_func(edges):
...@@ -62,7 +62,7 @@ def test_v2v_update_all(idtype): ...@@ -62,7 +62,7 @@ def test_v2v_update_all(idtype):
# test 2d node features # test 2d node features
_test('f2') _test('f2')
@parametrize_dtype @parametrize_idtype
def test_v2v_snr(idtype): def test_v2v_snr(idtype):
u = F.tensor([0, 0, 0, 3, 4, 9], idtype) u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
v = F.tensor([1, 2, 3, 9, 9, 0], idtype) v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
...@@ -106,7 +106,7 @@ def test_v2v_snr(idtype): ...@@ -106,7 +106,7 @@ def test_v2v_snr(idtype):
_test('f2') _test('f2')
@parametrize_dtype @parametrize_idtype
def test_v2v_pull(idtype): def test_v2v_pull(idtype):
nodes = F.tensor([1, 2, 3, 9], idtype) nodes = F.tensor([1, 2, 3, 9], idtype)
def _test(fld): def _test(fld):
...@@ -147,7 +147,7 @@ def test_v2v_pull(idtype): ...@@ -147,7 +147,7 @@ def test_v2v_pull(idtype):
# test 2d node features # test 2d node features
_test('f2') _test('f2')
@parametrize_dtype @parametrize_idtype
def test_update_all_multi_fallback(idtype): def test_update_all_multi_fallback(idtype):
# create a graph with zero in degree nodes # create a graph with zero in degree nodes
g = dgl.DGLGraph() g = dgl.DGLGraph()
...@@ -193,7 +193,7 @@ def test_update_all_multi_fallback(idtype): ...@@ -193,7 +193,7 @@ def test_update_all_multi_fallback(idtype):
_afunc) _afunc)
assert F.allclose(o2, g.ndata.pop('o2')) assert F.allclose(o2, g.ndata.pop('o2'))
@parametrize_dtype @parametrize_idtype
def test_pull_multi_fallback(idtype): def test_pull_multi_fallback(idtype):
# create a graph with zero in degree nodes # create a graph with zero in degree nodes
g = dgl.DGLGraph() g = dgl.DGLGraph()
...@@ -247,7 +247,7 @@ def test_pull_multi_fallback(idtype): ...@@ -247,7 +247,7 @@ def test_pull_multi_fallback(idtype):
nodes = [0, 1, 2, 9] nodes = [0, 1, 2, 9]
_pull_nodes(nodes) _pull_nodes(nodes)
@parametrize_dtype @parametrize_idtype
def test_spmv_3d_feat(idtype): def test_spmv_3d_feat(idtype):
def src_mul_edge_udf(edges): def src_mul_edge_udf(edges):
return {'sum': edges.src['h'] * F.unsqueeze(F.unsqueeze(edges.data['h'], 1), 1)} return {'sum': edges.src['h'] * F.unsqueeze(F.unsqueeze(edges.data['h'], 1), 1)}
......
tests/compute/test_subgraph.py
View file @ 14251504
...@@ -6,7 +6,7 @@ import pytest ...@@ -6,7 +6,7 @@ import pytest
import dgl import dgl
import backend as F import backend as F
from test_utils import parametrize_dtype from test_utils import parametrize_idtype
D = 5 D = 5
...@@ -118,7 +118,7 @@ def create_test_heterograph(idtype): ...@@ -118,7 +118,7 @@ def create_test_heterograph(idtype):
return g return g
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support bool tensor") @unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support bool tensor")
@parametrize_dtype @parametrize_idtype
def test_subgraph_mask(idtype): def test_subgraph_mask(idtype):
g = create_test_heterograph(idtype) g = create_test_heterograph(idtype)
g_graph = g['follows'] g_graph = g['follows']
...@@ -158,7 +158,7 @@ def test_subgraph_mask(idtype): ...@@ -158,7 +158,7 @@ def test_subgraph_mask(idtype):
'wishes': F.tensor([False, True], dtype=F.bool)}) 'wishes': F.tensor([False, True], dtype=F.bool)})
_check_subgraph(g, sg2) _check_subgraph(g, sg2)
@parametrize_dtype @parametrize_idtype
def test_subgraph1(idtype): def test_subgraph1(idtype):
g = create_test_heterograph(idtype) g = create_test_heterograph(idtype)
g_graph = g['follows'] g_graph = g['follows']
...@@ -310,7 +310,7 @@ def test_subgraph1(idtype): ...@@ -310,7 +310,7 @@ def test_subgraph1(idtype):
dst = F.asnumpy(dst) dst = F.asnumpy(dst)
assert np.array_equal(src, np.array([1])) assert np.array_equal(src, np.array([1]))
@parametrize_dtype @parametrize_idtype
def test_in_subgraph(idtype): def test_in_subgraph(idtype):
hg = dgl.heterograph({ hg = dgl.heterograph({
('user', 'follow', 'user'): ([1, 2, 3, 0, 2, 3, 0], [0, 0, 0, 1, 1, 1, 2]), ('user', 'follow', 'user'): ([1, 2, 3, 0, 2, 3, 0], [0, 0, 0, 1, 1, 1, 2]),
...@@ -377,7 +377,7 @@ def test_in_subgraph(idtype): ...@@ -377,7 +377,7 @@ def test_in_subgraph(idtype):
assert subg.num_nodes('coin') == 0 assert subg.num_nodes('coin') == 0
assert subg.num_edges('flips') == 0 assert subg.num_edges('flips') == 0
@parametrize_dtype @parametrize_idtype
def test_out_subgraph(idtype): def test_out_subgraph(idtype):
hg = dgl.heterograph({ hg = dgl.heterograph({
('user', 'follow', 'user'): ([1, 2, 3, 0, 2, 3, 0], [0, 0, 0, 1, 1, 1, 2]), ('user', 'follow', 'user'): ([1, 2, 3, 0, 2, 3, 0], [0, 0, 0, 1, 1, 1, 2]),
...@@ -459,7 +459,7 @@ def test_subgraph_message_passing(): ...@@ -459,7 +459,7 @@ def test_subgraph_message_passing():
sg = g.subgraph([1, 2, 3]).to(F.ctx()) sg = g.subgraph([1, 2, 3]).to(F.ctx())
sg.update_all(lambda edges: {'x': edges.src['x']}, lambda nodes: {'y': F.sum(nodes.mailbox['x'], 1)}) sg.update_all(lambda edges: {'x': edges.src['x']}, lambda nodes: {'y': F.sum(nodes.mailbox['x'], 1)})
@parametrize_dtype @parametrize_idtype
def test_khop_in_subgraph(idtype): def test_khop_in_subgraph(idtype):
g = dgl.graph(([1, 1, 2, 3, 4], [0, 2, 0, 4, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([1, 1, 2, 3, 4], [0, 2, 0, 4, 2]), idtype=idtype, device=F.ctx())
g.edata['w'] = F.tensor([ g.edata['w'] = F.tensor([
...@@ -535,7 +535,7 @@ def test_khop_in_subgraph(idtype): ...@@ -535,7 +535,7 @@ def test_khop_in_subgraph(idtype):
assert F.array_equal(F.astype(inv['user'], idtype), F.tensor([0, 1], idtype)) assert F.array_equal(F.astype(inv['user'], idtype), F.tensor([0, 1], idtype))
assert F.array_equal(F.astype(inv['game'], idtype), F.tensor([0], idtype)) assert F.array_equal(F.astype(inv['game'], idtype), F.tensor([0], idtype))
@parametrize_dtype @parametrize_idtype
def test_khop_out_subgraph(idtype): def test_khop_out_subgraph(idtype):
g = dgl.graph(([0, 2, 0, 4, 2], [1, 1, 2, 3, 4]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 2, 0, 4, 2], [1, 1, 2, 3, 4]), idtype=idtype, device=F.ctx())
g.edata['w'] = F.tensor([ g.edata['w'] = F.tensor([
...@@ -646,7 +646,7 @@ def test_subframes(parent_idx_device, child_device): ...@@ -646,7 +646,7 @@ def test_subframes(parent_idx_device, child_device):
@unittest.skipIf(F._default_context_str != "gpu", reason="UVA only available on GPU") @unittest.skipIf(F._default_context_str != "gpu", reason="UVA only available on GPU")
@pytest.mark.parametrize('device', [F.cpu(), F.cuda()]) @pytest.mark.parametrize('device', [F.cpu(), F.cuda()])
@unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="UVA only supported for PyTorch") @unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="UVA only supported for PyTorch")
@parametrize_dtype @parametrize_idtype
def test_uva_subgraph(idtype, device): def test_uva_subgraph(idtype, device):
g = create_test_heterograph(idtype) g = create_test_heterograph(idtype)
g = g.to(F.cpu()) g = g.to(F.cpu())
......
tests/compute/test_transform.py
View file @ 14251504
...@@ -26,7 +26,7 @@ import unittest ...@@ -26,7 +26,7 @@ import unittest
import math import math
import pytest import pytest
from test_utils.graph_cases import get_cases from test_utils.graph_cases import get_cases
from utils import parametrize_dtype from test_utils import parametrize_idtype
from test_heterograph import create_test_heterograph3, create_test_heterograph4, create_test_heterograph5 from test_heterograph import create_test_heterograph3, create_test_heterograph4, create_test_heterograph5
...@@ -43,7 +43,7 @@ def test_line_graph1(): ...@@ -43,7 +43,7 @@ def test_line_graph1():
assert F.allclose(L.ndata['h'], G.edata['h']) assert F.allclose(L.ndata['h'], G.edata['h'])
assert G.device == F.ctx() assert G.device == F.ctx()
@parametrize_dtype @parametrize_idtype
def test_line_graph2(idtype): def test_line_graph2(idtype):
g = dgl.heterograph({ g = dgl.heterograph({
('user', 'follows', 'user'): ([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]) ('user', 'follows', 'user'): ([0, 1, 1, 2, 2],[2, 0, 2, 0, 1])
...@@ -105,7 +105,7 @@ def test_no_backtracking(): ...@@ -105,7 +105,7 @@ def test_no_backtracking():
assert not L.has_edge_between(e2, e1) assert not L.has_edge_between(e2, e1)
# reverse graph related # reverse graph related
@parametrize_dtype @parametrize_idtype
def test_reverse(idtype): def test_reverse(idtype):
g = dgl.DGLGraph() g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -239,7 +239,7 @@ def test_reverse(idtype): ...@@ -239,7 +239,7 @@ def test_reverse(idtype):
assert ('hhh' in g_r.edges['follows'].data) is True assert ('hhh' in g_r.edges['follows'].data) is True
@parametrize_dtype @parametrize_idtype
def test_reverse_shared_frames(idtype): def test_reverse_shared_frames(idtype):
g = dgl.DGLGraph() g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -543,7 +543,7 @@ def test_partition_with_halo(): ...@@ -543,7 +543,7 @@ def test_partition_with_halo():
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet') @unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F._default_context_str == 'gpu', reason="METIS doesn't support GPU") @unittest.skipIf(F._default_context_str == 'gpu', reason="METIS doesn't support GPU")
@parametrize_dtype @parametrize_idtype
def test_metis_partition(idtype): def test_metis_partition(idtype):
# TODO(zhengda) Metis fails to partition a small graph. # TODO(zhengda) Metis fails to partition a small graph.
g = create_large_graph(1000, idtype=idtype) g = create_large_graph(1000, idtype=idtype)
...@@ -675,7 +675,7 @@ def test_reorder_nodes(): ...@@ -675,7 +675,7 @@ def test_reorder_nodes():
old_neighs2 = g.predecessors(old_nid) old_neighs2 = g.predecessors(old_nid)
assert np.all(np.sort(old_neighs1) == np.sort(F.asnumpy(old_neighs2))) assert np.all(np.sort(old_neighs1) == np.sort(F.asnumpy(old_neighs2)))
@parametrize_dtype @parametrize_idtype
def test_compact(idtype): def test_compact(idtype):
g1 = dgl.heterograph({ g1 = dgl.heterograph({
('user', 'follow', 'user'): ([1, 3], [3, 5]), ('user', 'follow', 'user'): ([1, 3], [3, 5]),
...@@ -776,7 +776,7 @@ def test_compact(idtype): ...@@ -776,7 +776,7 @@ def test_compact(idtype):
_check(g4, new_g4, induced_nodes) _check(g4, new_g4, induced_nodes)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU to simple not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU to simple not implemented")
@parametrize_dtype @parametrize_idtype
def test_to_simple(idtype): def test_to_simple(idtype):
# homogeneous graph # homogeneous graph
g = dgl.graph((F.tensor([0, 1, 2, 1]), F.tensor([1, 2, 0, 2]))) g = dgl.graph((F.tensor([0, 1, 2, 1]), F.tensor([1, 2, 0, 2])))
...@@ -877,7 +877,7 @@ def test_to_simple(idtype): ...@@ -877,7 +877,7 @@ def test_to_simple(idtype):
sg = dgl.to_simple(g) sg = dgl.to_simple(g)
assert F.array_equal(sg.edge_ids(u, v), eids) assert F.array_equal(sg.edge_ids(u, v), eids)
@parametrize_dtype @parametrize_idtype
def test_to_block(idtype): def test_to_block(idtype):
def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True): def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True):
if dst_nodes is not None: if dst_nodes is not None:
...@@ -1008,7 +1008,7 @@ def test_to_block(idtype): ...@@ -1008,7 +1008,7 @@ def test_to_block(idtype):
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype @parametrize_idtype
def test_remove_edges(idtype): def test_remove_edges(idtype):
def check(g1, etype, g, edges_removed): def check(g1, etype, g, edges_removed):
src, dst, eid = g.edges(etype=etype, form='all') src, dst, eid = g.edges(etype=etype, form='all')
...@@ -1061,7 +1061,7 @@ def test_remove_edges(idtype): ...@@ -1061,7 +1061,7 @@ def test_remove_edges(idtype):
check(g4, 'AB', g, [3, 1, 2, 0]) check(g4, 'AB', g, [3, 1, 2, 0])
check(g4, 'BA', g, []) check(g4, 'BA', g, [])
@parametrize_dtype @parametrize_idtype
def test_add_edges(idtype): def test_add_edges(idtype):
# homogeneous graph # homogeneous graph
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
...@@ -1239,7 +1239,7 @@ def test_add_edges(idtype): ...@@ -1239,7 +1239,7 @@ def test_add_edges(idtype):
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 1, 1], dtype=idtype)) assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 1, 1], dtype=idtype))
assert F.array_equal(g.edges['develops'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype)) assert F.array_equal(g.edges['develops'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype))
@parametrize_dtype @parametrize_idtype
def test_add_nodes(idtype): def test_add_nodes(idtype):
# homogeneous Graphs # homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
...@@ -1277,7 +1277,7 @@ def test_add_nodes(idtype): ...@@ -1277,7 +1277,7 @@ def test_add_nodes(idtype):
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1, 0], dtype=idtype)) assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 2], dtype=idtype)) assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 2], dtype=idtype))
@parametrize_dtype @parametrize_idtype
def test_remove_edges(idtype): def test_remove_edges(idtype):
# homogeneous Graphs # homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
...@@ -1447,7 +1447,7 @@ def test_remove_edges(idtype): ...@@ -1447,7 +1447,7 @@ def test_remove_edges(idtype):
assert F.array_equal(bg.batch_num_edges('follows'), bg_r.batch_num_edges('follows')) assert F.array_equal(bg.batch_num_edges('follows'), bg_r.batch_num_edges('follows'))
assert F.array_equal(bg_r.batch_num_edges('plays'), F.tensor([1, 0, 1], dtype=F.int64)) assert F.array_equal(bg_r.batch_num_edges('plays'), F.tensor([1, 0, 1], dtype=F.int64))
@parametrize_dtype @parametrize_idtype
def test_remove_nodes(idtype): def test_remove_nodes(idtype):
# homogeneous Graphs # homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
...@@ -1622,7 +1622,7 @@ def test_remove_nodes(idtype): ...@@ -1622,7 +1622,7 @@ def test_remove_nodes(idtype):
assert F.array_equal(bg.batch_num_edges('follows'), bg_r.batch_num_edges('follows')) assert F.array_equal(bg.batch_num_edges('follows'), bg_r.batch_num_edges('follows'))
assert F.array_equal(bg_r.batch_num_edges('plays'), F.tensor([1, 0, 1], dtype=F.int64)) assert F.array_equal(bg_r.batch_num_edges('plays'), F.tensor([1, 0, 1], dtype=F.int64))
@parametrize_dtype @parametrize_idtype
def test_add_selfloop(idtype): def test_add_selfloop(idtype):
# homogeneous graph # homogeneous graph
g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
...@@ -1666,7 +1666,7 @@ def test_add_selfloop(idtype): ...@@ -1666,7 +1666,7 @@ def test_add_selfloop(idtype):
raise_error = True raise_error = True
assert raise_error assert raise_error
@parametrize_dtype @parametrize_idtype
def test_remove_selfloop(idtype): def test_remove_selfloop(idtype):
# homogeneous graph # homogeneous graph
g = dgl.graph(([0, 0, 0, 1], [1, 0, 0, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 0, 0, 1], [1, 0, 0, 2]), idtype=idtype, device=F.ctx())
...@@ -1707,7 +1707,7 @@ def test_remove_selfloop(idtype): ...@@ -1707,7 +1707,7 @@ def test_remove_selfloop(idtype):
assert raise_error assert raise_error
@parametrize_dtype @parametrize_idtype
def test_reorder_graph(idtype): def test_reorder_graph(idtype):
g = dgl.graph(([0, 1, 2, 3, 4], [2, 2, 3, 2, 3]), g = dgl.graph(([0, 1, 2, 3, 4], [2, 2, 3, 2, 3]),
idtype=idtype, device=F.ctx()) idtype=idtype, device=F.ctx())
...@@ -1821,7 +1821,7 @@ def test_reorder_graph(idtype): ...@@ -1821,7 +1821,7 @@ def test_reorder_graph(idtype):
#assert 'csc' in sum(rfg.formats().values(), []) #assert 'csc' in sum(rfg.formats().values(), [])
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support a slicing operation") @unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support a slicing operation")
@parametrize_dtype @parametrize_idtype
def test_norm_by_dst(idtype): def test_norm_by_dst(idtype):
# Case1: A homogeneous graph # Case1: A homogeneous graph
g = dgl.graph(([0, 1, 1], [1, 1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1, 1], [1, 1, 2]), idtype=idtype, device=F.ctx())
...@@ -1836,7 +1836,7 @@ def test_norm_by_dst(idtype): ...@@ -1836,7 +1836,7 @@ def test_norm_by_dst(idtype):
eweight = dgl.norm_by_dst(g, etype=('user', 'plays', 'game')) eweight = dgl.norm_by_dst(g, etype=('user', 'plays', 'game'))
assert F.allclose(eweight, F.tensor([0.5, 0.5, 1.0])) assert F.allclose(eweight, F.tensor([0.5, 0.5, 1.0]))
@parametrize_dtype @parametrize_idtype
def test_module_add_self_loop(idtype): def test_module_add_self_loop(idtype):
g = dgl.graph(([1, 1], [1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([1, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.randn((g.num_nodes(), 2)) g.ndata['h'] = F.randn((g.num_nodes(), 2))
...@@ -1914,7 +1914,7 @@ def test_module_add_self_loop(idtype): ...@@ -1914,7 +1914,7 @@ def test_module_add_self_loop(idtype):
assert 'w1' in new_g.edges['plays'].data assert 'w1' in new_g.edges['plays'].data
assert 'w2' in new_g.edges['follows'].data assert 'w2' in new_g.edges['follows'].data
@parametrize_dtype @parametrize_idtype
def test_module_remove_self_loop(idtype): def test_module_remove_self_loop(idtype):
transform = dgl.RemoveSelfLoop() transform = dgl.RemoveSelfLoop()
...@@ -1957,7 +1957,7 @@ def test_module_remove_self_loop(idtype): ...@@ -1957,7 +1957,7 @@ def test_module_remove_self_loop(idtype):
assert 'w1' in new_g.edges['plays'].data assert 'w1' in new_g.edges['plays'].data
assert 'w2' in new_g.edges['follows'].data assert 'w2' in new_g.edges['follows'].data
@parametrize_dtype @parametrize_idtype
def test_module_add_reverse(idtype): def test_module_add_reverse(idtype):
transform = dgl.AddReverse() transform = dgl.AddReverse()
...@@ -2044,7 +2044,7 @@ def test_module_add_reverse(idtype): ...@@ -2044,7 +2044,7 @@ def test_module_add_reverse(idtype):
assert eset == {(2, 1), (2, 2)} assert eset == {(2, 1), (2, 2)}
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not supported for to_simple") @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not supported for to_simple")
@parametrize_dtype @parametrize_idtype
def test_module_to_simple(idtype): def test_module_to_simple(idtype):
transform = dgl.ToSimple() transform = dgl.ToSimple()
g = dgl.graph(([0, 1, 1], [1, 2, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1, 1], [1, 2, 2]), idtype=idtype, device=F.ctx())
...@@ -2083,7 +2083,7 @@ def test_module_to_simple(idtype): ...@@ -2083,7 +2083,7 @@ def test_module_to_simple(idtype):
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == {(0, 1), (1, 1)} assert eset == {(0, 1), (1, 1)}
@parametrize_dtype @parametrize_idtype
def test_module_line_graph(idtype): def test_module_line_graph(idtype):
transform = dgl.LineGraph() transform = dgl.LineGraph()
g = dgl.graph(([0, 1, 1], [1, 0, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1, 1], [1, 0, 2]), idtype=idtype, device=F.ctx())
...@@ -2106,7 +2106,7 @@ def test_module_line_graph(idtype): ...@@ -2106,7 +2106,7 @@ def test_module_line_graph(idtype):
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == {(0, 2)} assert eset == {(0, 2)}
@parametrize_dtype @parametrize_idtype
def test_module_khop_graph(idtype): def test_module_khop_graph(idtype):
transform = dgl.KHopGraph(2) transform = dgl.KHopGraph(2)
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
...@@ -2120,7 +2120,7 @@ def test_module_khop_graph(idtype): ...@@ -2120,7 +2120,7 @@ def test_module_khop_graph(idtype):
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == {(0, 2)} assert eset == {(0, 2)}
@parametrize_dtype @parametrize_idtype
def test_module_add_metapaths(idtype): def test_module_add_metapaths(idtype):
g = dgl.heterograph({ g = dgl.heterograph({
('person', 'author', 'paper'): ([0, 0, 1], [1, 2, 2]), ('person', 'author', 'paper'): ([0, 0, 1], [1, 2, 2]),
...@@ -2179,7 +2179,7 @@ def test_module_add_metapaths(idtype): ...@@ -2179,7 +2179,7 @@ def test_module_add_metapaths(idtype):
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == {(0, 1), (1, 1)} assert eset == {(0, 1), (1, 1)}
@parametrize_dtype @parametrize_idtype
def test_module_compose(idtype): def test_module_compose(idtype):
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
transform = dgl.Compose([dgl.AddReverse(), dgl.AddSelfLoop()]) transform = dgl.Compose([dgl.AddReverse(), dgl.AddSelfLoop()])
...@@ -2192,7 +2192,7 @@ def test_module_compose(idtype): ...@@ -2192,7 +2192,7 @@ def test_module_compose(idtype):
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == {(0, 1), (1, 2), (1, 0), (2, 1), (0, 0), (1, 1), (2, 2)} assert eset == {(0, 1), (1, 2), (1, 0), (2, 1), (0, 0), (1, 1), (2, 2)}
@parametrize_dtype @parametrize_idtype
def test_module_gcnnorm(idtype): def test_module_gcnnorm(idtype):
g = dgl.heterograph({ g = dgl.heterograph({
('A', 'r1', 'A'): ([0, 1, 2], [0, 0, 1]), ('A', 'r1', 'A'): ([0, 1, 2], [0, 0, 1]),
...@@ -2208,7 +2208,7 @@ def test_module_gcnnorm(idtype): ...@@ -2208,7 +2208,7 @@ def test_module_gcnnorm(idtype):
assert F.allclose(new_g.edges[('B', 'r3', 'B')].data['w'], F.tensor([1./3, 2./3, 0.])) assert F.allclose(new_g.edges[('B', 'r3', 'B')].data['w'], F.tensor([1./3, 2./3, 0.]))
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
def test_module_ppr(idtype): def test_module_ppr(idtype):
g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.randn((6, 2)) g.ndata['h'] = F.randn((6, 2))
...@@ -2233,7 +2233,7 @@ def test_module_ppr(idtype): ...@@ -2233,7 +2233,7 @@ def test_module_ppr(idtype):
(3, 3), (3, 5), (4, 3), (4, 4), (4, 5), (5, 5)} (3, 3), (3, 5), (4, 3), (4, 4), (4, 5), (5, 5)}
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
def test_module_heat_kernel(idtype): def test_module_heat_kernel(idtype):
# Case1: directed graph # Case1: directed graph
g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype, device=F.ctx())
...@@ -2255,7 +2255,7 @@ def test_module_heat_kernel(idtype): ...@@ -2255,7 +2255,7 @@ def test_module_heat_kernel(idtype):
assert eset == {(0, 0), (1, 1), (2, 2), (3, 3)} assert eset == {(0, 0), (1, 1), (2, 2), (3, 3)}
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
def test_module_gdc(idtype): def test_module_gdc(idtype):
transform = dgl.GDC([0.1, 0.2, 0.1], avg_degree=1) transform = dgl.GDC([0.1, 0.2, 0.1], avg_degree=1)
g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype, device=F.ctx())
...@@ -2278,7 +2278,7 @@ def test_module_gdc(idtype): ...@@ -2278,7 +2278,7 @@ def test_module_gdc(idtype):
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == {(0, 0), (1, 1), (2, 2), (3, 3), (4, 3), (4, 4), (5, 5)} assert eset == {(0, 0), (1, 1), (2, 2), (3, 3), (4, 3), (4, 4), (5, 5)}
@parametrize_dtype @parametrize_idtype
def test_module_node_shuffle(idtype): def test_module_node_shuffle(idtype):
transform = dgl.NodeShuffle() transform = dgl.NodeShuffle()
g = dgl.heterograph({ g = dgl.heterograph({
...@@ -2287,7 +2287,7 @@ def test_module_node_shuffle(idtype): ...@@ -2287,7 +2287,7 @@ def test_module_node_shuffle(idtype):
new_g = transform(g) new_g = transform(g)
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
def test_module_drop_node(idtype): def test_module_drop_node(idtype):
transform = dgl.DropNode() transform = dgl.DropNode()
g = dgl.heterograph({ g = dgl.heterograph({
...@@ -2300,7 +2300,7 @@ def test_module_drop_node(idtype): ...@@ -2300,7 +2300,7 @@ def test_module_drop_node(idtype):
assert new_g.canonical_etypes == g.canonical_etypes assert new_g.canonical_etypes == g.canonical_etypes
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
def test_module_drop_edge(idtype): def test_module_drop_edge(idtype):
transform = dgl.DropEdge() transform = dgl.DropEdge()
g = dgl.heterograph({ g = dgl.heterograph({
...@@ -2313,7 +2313,7 @@ def test_module_drop_edge(idtype): ...@@ -2313,7 +2313,7 @@ def test_module_drop_edge(idtype):
assert new_g.ntypes == g.ntypes assert new_g.ntypes == g.ntypes
assert new_g.canonical_etypes == g.canonical_etypes assert new_g.canonical_etypes == g.canonical_etypes
@parametrize_dtype @parametrize_idtype
def test_module_add_edge(idtype): def test_module_add_edge(idtype):
transform = dgl.AddEdge() transform = dgl.AddEdge()
g = dgl.heterograph({ g = dgl.heterograph({
...@@ -2328,7 +2328,7 @@ def test_module_add_edge(idtype): ...@@ -2328,7 +2328,7 @@ def test_module_add_edge(idtype):
assert new_g.ntypes == g.ntypes assert new_g.ntypes == g.ntypes
assert new_g.canonical_etypes == g.canonical_etypes assert new_g.canonical_etypes == g.canonical_etypes
@parametrize_dtype @parametrize_idtype
def test_module_random_walk_pe(idtype): def test_module_random_walk_pe(idtype):
transform = dgl.RandomWalkPE(2, 'rwpe') transform = dgl.RandomWalkPE(2, 'rwpe')
g = dgl.graph(([0, 1, 1], [1, 1, 0]), idtype=idtype, device=F.ctx()) g = dgl.graph(([0, 1, 1], [1, 1, 0]), idtype=idtype, device=F.ctx())
...@@ -2336,7 +2336,7 @@ def test_module_random_walk_pe(idtype): ...@@ -2336,7 +2336,7 @@ def test_module_random_walk_pe(idtype):
tgt = F.copy_to(F.tensor([[0., 0.5],[0.5, 0.75]]), g.device) tgt = F.copy_to(F.tensor([[0., 0.5],[0.5, 0.75]]), g.device)
assert F.allclose(new_g.ndata['rwpe'], tgt) assert F.allclose(new_g.ndata['rwpe'], tgt)
@parametrize_dtype @parametrize_idtype
def test_module_laplacian_pe(idtype): def test_module_laplacian_pe(idtype):
transform = dgl.LaplacianPE(2, 'lappe') transform = dgl.LaplacianPE(2, 'lappe')
g = dgl.graph(([2, 1, 0, 3, 1, 1],[3, 0, 1, 3, 3, 1]), idtype=idtype, device=F.ctx()) g = dgl.graph(([2, 1, 0, 3, 1, 1],[3, 0, 1, 3, 3, 1]), idtype=idtype, device=F.ctx())
...@@ -2422,7 +2422,7 @@ def test_module_sign(g): ...@@ -2422,7 +2422,7 @@ def test_module_sign(g):
assert torch.allclose(g.ndata['out_feat_1'], target) assert torch.allclose(g.ndata['out_feat_1'], target)
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
def test_module_row_feat_normalizer(idtype): def test_module_row_feat_normalizer(idtype):
# Case1: Normalize features of a homogeneous graph. # Case1: Normalize features of a homogeneous graph.
transform = dgl.RowFeatNormalizer(subtract_min=True) transform = dgl.RowFeatNormalizer(subtract_min=True)
...@@ -2457,7 +2457,7 @@ def test_module_row_feat_normalizer(idtype): ...@@ -2457,7 +2457,7 @@ def test_module_row_feat_normalizer(idtype):
assert F.allclose(g.edata['w'][('player', 'plays', 'game')].sum(1), F.tensor([1.0, 1.0])) assert F.allclose(g.edata['w'][('player', 'plays', 'game')].sum(1), F.tensor([1.0, 1.0]))
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now') @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@parametrize_dtype @parametrize_idtype
def test_module_feat_mask(idtype): def test_module_feat_mask(idtype):
# Case1: Mask node and edge feature tensors of a homogeneous graph. # Case1: Mask node and edge feature tensors of a homogeneous graph.
transform = dgl.FeatMask() transform = dgl.FeatMask()
......
tests/compute/test_traversal.py
View file @ 14251504
...@@ -10,7 +10,7 @@ import scipy.sparse as sp ...@@ -10,7 +10,7 @@ import scipy.sparse as sp
import backend as F import backend as F
import itertools import itertools
from utils import parametrize_dtype from test_utils import parametrize_idtype
np.random.seed(42) np.random.seed(42)
...@@ -18,7 +18,7 @@ def toset(x): ...@@ -18,7 +18,7 @@ def toset(x):
# F.zerocopy_to_numpy may return a int # F.zerocopy_to_numpy may return a int
return set(F.zerocopy_to_numpy(x).tolist()) return set(F.zerocopy_to_numpy(x).tolist())
@parametrize_dtype @parametrize_idtype
def test_bfs(idtype, n=100): def test_bfs(idtype, n=100):
def _bfs_nx(g_nx, src): def _bfs_nx(g_nx, src):
edges = nx.bfs_edges(g_nx, src) edges = nx.bfs_edges(g_nx, src)
...@@ -59,7 +59,7 @@ def test_bfs(idtype, n=100): ...@@ -59,7 +59,7 @@ def test_bfs(idtype, n=100):
assert len(edges_dgl) == len(edges_nx) assert len(edges_dgl) == len(edges_nx)
assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx)) assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))
@parametrize_dtype @parametrize_idtype
def test_topological_nodes(idtype, n=100): def test_topological_nodes(idtype, n=100):
a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n)) a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
b = sp.tril(a, -1).tocoo() b = sp.tril(a, -1).tocoo()
...@@ -86,7 +86,7 @@ def test_topological_nodes(idtype, n=100): ...@@ -86,7 +86,7 @@ def test_topological_nodes(idtype, n=100):
assert all(toset(x) == toset(y) for x, y in zip(layers_dgl, layers_spmv)) assert all(toset(x) == toset(y) for x, y in zip(layers_dgl, layers_spmv))
DFS_LABEL_NAMES = ['forward', 'reverse', 'nontree'] DFS_LABEL_NAMES = ['forward', 'reverse', 'nontree']
@parametrize_dtype @parametrize_idtype
def test_dfs_labeled_edges(idtype, example=False): def test_dfs_labeled_edges(idtype, example=False):
dgl_g = dgl.DGLGraph().astype(idtype) dgl_g = dgl.DGLGraph().astype(idtype)
dgl_g.add_nodes(6) dgl_g.add_nodes(6)
......
tests/compute/utils.py
View file @ 14251504
...@@ -3,12 +3,6 @@ import backend as F ...@@ -3,12 +3,6 @@ import backend as F
import dgl import dgl
from dgl.base import is_internal_column from dgl.base import is_internal_column
if F._default_context_str == 'cpu':
parametrize_dtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
else:
# only test int32 on GPU because many graph operators are not supported for int64.
parametrize_dtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
def check_fail(fn, *args, **kwargs): def check_fail(fn, *args, **kwargs):
try: try:
fn(*args, **kwargs) fn(*args, **kwargs)
......
tests/mxnet/test_nn.py
View file @ 14251504
...@@ -8,7 +8,7 @@ import dgl.nn.mxnet as nn ...@@ -8,7 +8,7 @@ import dgl.nn.mxnet as nn
import dgl.function as fn import dgl.function as fn
import backend as F import backend as F
from test_utils.graph_cases import get_cases, random_graph, random_bipartite, random_dglgraph from test_utils.graph_cases import get_cases, random_graph, random_bipartite, random_dglgraph
from test_utils import parametrize_dtype from test_utils import parametrize_idtype
from mxnet import autograd, gluon, nd from mxnet import autograd, gluon, nd
def check_close(a, b): def check_close(a, b):
...@@ -19,7 +19,7 @@ def _AXWb(A, X, W, b): ...@@ -19,7 +19,7 @@ def _AXWb(A, X, W, b):
Y = mx.nd.dot(A, X.reshape(X.shape[0], -1)).reshape(X.shape) Y = mx.nd.dot(A, X.reshape(X.shape[0], -1)).reshape(X.shape)
return Y + b.data(X.context) return Y + b.data(X.context)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_graph_conv(idtype, out_dim): def test_graph_conv(idtype, out_dim):
g = dgl.from_networkx(nx.path_graph(3)) g = dgl.from_networkx(nx.path_graph(3))
...@@ -79,7 +79,7 @@ def test_graph_conv(idtype, out_dim): ...@@ -79,7 +79,7 @@ def test_graph_conv(idtype, out_dim):
assert "h" in g.ndata assert "h" in g.ndata
check_close(g.ndata['h'], 2 * F.ones((3, 1))) check_close(g.ndata['h'], 2 * F.ones((3, 1)))
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right', 'left']) @pytest.mark.parametrize('norm', ['none', 'both', 'right', 'left'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -99,7 +99,7 @@ def test_graph_conv2(idtype, g, norm, weight, bias, out_dim): ...@@ -99,7 +99,7 @@ def test_graph_conv2(idtype, g, norm, weight, bias, out_dim):
h_out = conv(g, h, ext_w) h_out = conv(g, h, ext_w)
assert h_out.shape == (ndst, out_dim) assert h_out.shape == (ndst, out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right']) @pytest.mark.parametrize('norm', ['none', 'both', 'right'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -161,7 +161,7 @@ def test_tagconv(out_dim): ...@@ -161,7 +161,7 @@ def test_tagconv(out_dim):
h1 = conv(g, h0) h1 = conv(g, h0)
assert h1.shape[-1] == out_dim assert h1.shape[-1] == out_dim
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 20]) @pytest.mark.parametrize('out_dim', [1, 20])
@pytest.mark.parametrize('num_heads', [1, 5]) @pytest.mark.parametrize('num_heads', [1, 5])
...@@ -182,7 +182,7 @@ def test_gat_conv(g, idtype, out_dim, num_heads): ...@@ -182,7 +182,7 @@ def test_gat_conv(g, idtype, out_dim, num_heads):
gat.initialize(ctx=ctx) gat.initialize(ctx=ctx)
h = gat(g, feat) h = gat(g, feat)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -197,7 +197,7 @@ def test_gat_conv_bi(g, idtype, out_dim, num_heads): ...@@ -197,7 +197,7 @@ def test_gat_conv_bi(g, idtype, out_dim, num_heads):
_, a = gat(g, feat, True) _, a = gat(g, feat, True)
assert a.shape == (g.number_of_edges(), num_heads, 1) assert a.shape == (g.number_of_edges(), num_heads, 1)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite']))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn'])
@pytest.mark.parametrize('out_dim', [1, 10]) @pytest.mark.parametrize('out_dim', [1, 10])
...@@ -210,7 +210,7 @@ def test_sage_conv(idtype, g, aggre_type, out_dim): ...@@ -210,7 +210,7 @@ def test_sage_conv(idtype, g, aggre_type, out_dim):
h = sage(g, feat) h = sage(g, feat)
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize('g', get_cases(['bipartite']))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn'])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
...@@ -225,7 +225,7 @@ def test_sage_conv_bi(idtype, g, aggre_type, out_dim): ...@@ -225,7 +225,7 @@ def test_sage_conv_bi(idtype, g, aggre_type, out_dim):
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == g.number_of_dst_nodes() assert h.shape[0] == g.number_of_dst_nodes()
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn'])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_sage_conv_bi2(idtype, aggre_type, out_dim): def test_sage_conv_bi2(idtype, aggre_type, out_dim):
...@@ -275,7 +275,7 @@ def test_cheb_conv(out_dim): ...@@ -275,7 +275,7 @@ def test_cheb_conv(out_dim):
h1 = cheb(g, h0) h1 = cheb(g, h0)
assert h1.shape == (20, out_dim) assert h1.shape == (20, out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
def test_agnn_conv(g, idtype): def test_agnn_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -287,7 +287,7 @@ def test_agnn_conv(g, idtype): ...@@ -287,7 +287,7 @@ def test_agnn_conv(g, idtype):
h = agnn_conv(g, feat) h = agnn_conv(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 10) assert h.shape == (g.number_of_dst_nodes(), 10)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
def test_agnn_conv_bi(g, idtype): def test_agnn_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -335,7 +335,7 @@ def test_dense_cheb_conv(out_dim): ...@@ -335,7 +335,7 @@ def test_dense_cheb_conv(out_dim):
out_dense_cheb = dense_cheb(adj, feat, 2.0) out_dense_cheb = dense_cheb(adj, feat, 2.0)
assert F.allclose(out_cheb, out_dense_cheb) assert F.allclose(out_cheb, out_dense_cheb)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('norm_type', ['both', 'right', 'none']) @pytest.mark.parametrize('norm_type', ['both', 'right', 'none'])
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
...@@ -356,7 +356,7 @@ def test_dense_graph_conv(idtype, g, norm_type, out_dim): ...@@ -356,7 +356,7 @@ def test_dense_graph_conv(idtype, g, norm_type, out_dim):
out_dense_conv = dense_conv(adj, feat) out_dense_conv = dense_conv(adj, feat)
assert F.allclose(out_conv, out_dense_conv) assert F.allclose(out_conv, out_dense_conv)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'bipartite', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'bipartite', 'block-bipartite']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_dense_sage_conv(idtype, g, out_dim): def test_dense_sage_conv(idtype, g, out_dim):
...@@ -383,7 +383,7 @@ def test_dense_sage_conv(idtype, g, out_dim): ...@@ -383,7 +383,7 @@ def test_dense_sage_conv(idtype, g, out_dim):
out_dense_sage = dense_sage(adj, feat) out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage) assert F.allclose(out_sage, out_dense_sage)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_edge_conv(g, idtype, out_dim): def test_edge_conv(g, idtype, out_dim):
...@@ -397,7 +397,7 @@ def test_edge_conv(g, idtype, out_dim): ...@@ -397,7 +397,7 @@ def test_edge_conv(g, idtype, out_dim):
h1 = edge_conv(g, h0) h1 = edge_conv(g, h0)
assert h1.shape == (g.number_of_dst_nodes(), out_dim) assert h1.shape == (g.number_of_dst_nodes(), out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_edge_conv_bi(g, idtype, out_dim): def test_edge_conv_bi(g, idtype, out_dim):
...@@ -412,7 +412,7 @@ def test_edge_conv_bi(g, idtype, out_dim): ...@@ -412,7 +412,7 @@ def test_edge_conv_bi(g, idtype, out_dim):
h1 = edge_conv(g, (h0, x0)) h1 = edge_conv(g, (h0, x0))
assert h1.shape == (g.number_of_dst_nodes(), out_dim) assert h1.shape == (g.number_of_dst_nodes(), out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite']))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum'])
def test_gin_conv(g, idtype, aggregator_type): def test_gin_conv(g, idtype, aggregator_type):
...@@ -428,7 +428,7 @@ def test_gin_conv(g, idtype, aggregator_type): ...@@ -428,7 +428,7 @@ def test_gin_conv(g, idtype, aggregator_type):
h = gin_conv(g, feat) h = gin_conv(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 5) assert h.shape == (g.number_of_dst_nodes(), 5)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize('g', get_cases(['bipartite']))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum'])
def test_gin_conv_bi(g, idtype, aggregator_type): def test_gin_conv_bi(g, idtype, aggregator_type):
...@@ -445,7 +445,7 @@ def test_gin_conv_bi(g, idtype, aggregator_type): ...@@ -445,7 +445,7 @@ def test_gin_conv_bi(g, idtype, aggregator_type):
return h.shape == (g.number_of_dst_nodes(), 5) return h.shape == (g.number_of_dst_nodes(), 5)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
def test_gmm_conv(g, idtype): def test_gmm_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -457,7 +457,7 @@ def test_gmm_conv(g, idtype): ...@@ -457,7 +457,7 @@ def test_gmm_conv(g, idtype):
h1 = gmm_conv(g, h0, pseudo) h1 = gmm_conv(g, h0, pseudo)
assert h1.shape == (g.number_of_dst_nodes(), 2) assert h1.shape == (g.number_of_dst_nodes(), 2)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
def test_gmm_conv_bi(g, idtype): def test_gmm_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -471,7 +471,7 @@ def test_gmm_conv_bi(g, idtype): ...@@ -471,7 +471,7 @@ def test_gmm_conv_bi(g, idtype):
h1 = gmm_conv(g, (h0, hd), pseudo) h1 = gmm_conv(g, (h0, hd), pseudo)
assert h1.shape == (g.number_of_dst_nodes(), 2) assert h1.shape == (g.number_of_dst_nodes(), 2)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite']))
def test_nn_conv(g, idtype): def test_nn_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -484,7 +484,7 @@ def test_nn_conv(g, idtype): ...@@ -484,7 +484,7 @@ def test_nn_conv(g, idtype):
h1 = nn_conv(g, h0, etypes) h1 = nn_conv(g, h0, etypes)
assert h1.shape == (g.number_of_dst_nodes(), 2) assert h1.shape == (g.number_of_dst_nodes(), 2)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize('g', get_cases(['bipartite']))
def test_nn_conv_bi(g, idtype): def test_nn_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -715,7 +715,7 @@ def myagg(alist, dsttype): ...@@ -715,7 +715,7 @@ def myagg(alist, dsttype):
rst = rst + (i + 1) * alist[i] rst = rst + (i + 1) * alist[i]
return rst return rst
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg]) @pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg])
def test_hetero_conv(agg, idtype): def test_hetero_conv(agg, idtype):
g = dgl.heterograph({ g = dgl.heterograph({
......
tests/pytorch/test_dataloader.py
View file @ 14251504
...@@ -10,7 +10,7 @@ from torch.utils.data import DataLoader ...@@ -10,7 +10,7 @@ from torch.utils.data import DataLoader
from collections import defaultdict from collections import defaultdict
from collections.abc import Iterator, Mapping from collections.abc import Iterator, Mapping
from itertools import product from itertools import product
from test_utils import parametrize_dtype from test_utils import parametrize_idtype
import pytest import pytest
...@@ -128,7 +128,7 @@ def _check_device(data): ...@@ -128,7 +128,7 @@ def _check_device(data):
else: else:
assert data.device == F.ctx() assert data.device == F.ctx()
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('sampler_name', ['full', 'neighbor', 'neighbor2']) @pytest.mark.parametrize('sampler_name', ['full', 'neighbor', 'neighbor2'])
@pytest.mark.parametrize('pin_graph', [None, 'cuda_indices', 'cpu_indices']) @pytest.mark.parametrize('pin_graph', [None, 'cuda_indices', 'cpu_indices'])
def test_node_dataloader(idtype, sampler_name, pin_graph): def test_node_dataloader(idtype, sampler_name, pin_graph):
......
tests/pytorch/test_geometry.py
View file @ 14251504
...@@ -7,7 +7,7 @@ import torch as th ...@@ -7,7 +7,7 @@ import torch as th
from dgl import DGLError from dgl import DGLError
from dgl.base import DGLWarning from dgl.base import DGLWarning
from dgl.geometry import neighbor_matching, farthest_point_sampler from dgl.geometry import neighbor_matching, farthest_point_sampler
from test_utils import parametrize_dtype from test_utils import parametrize_idtype
from test_utils.graph_cases import get_cases from test_utils.graph_cases import get_cases
...@@ -181,7 +181,7 @@ def test_knn_cuda(algorithm, dist): ...@@ -181,7 +181,7 @@ def test_knn_cuda(algorithm, dist):
g = kg(x_empty, [3, 5], algorithm, dist) g = kg(x_empty, [3, 5], algorithm, dist)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph']))
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
@pytest.mark.parametrize('relabel', [True, False]) @pytest.mark.parametrize('relabel', [True, False])
......
tests/pytorch/test_nn.py
View file @ 14251504
...@@ -7,7 +7,7 @@ import dgl.function as fn ...@@ -7,7 +7,7 @@ import dgl.function as fn
import backend as F import backend as F
import pytest import pytest
from test_utils.graph_cases import get_cases, random_graph, random_bipartite, random_dglgraph from test_utils.graph_cases import get_cases, random_graph, random_bipartite, random_dglgraph
from test_utils import parametrize_dtype from test_utils import parametrize_idtype
from copy import deepcopy from copy import deepcopy
import pickle import pickle
...@@ -79,7 +79,7 @@ def test_graph_conv0(out_dim): ...@@ -79,7 +79,7 @@ def test_graph_conv0(out_dim):
new_weight = conv.weight.data new_weight = conv.weight.data
assert not F.allclose(old_weight, new_weight) assert not F.allclose(old_weight, new_weight)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right', 'left']) @pytest.mark.parametrize('norm', ['none', 'both', 'right', 'left'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -99,7 +99,7 @@ def test_graph_conv(idtype, g, norm, weight, bias, out_dim): ...@@ -99,7 +99,7 @@ def test_graph_conv(idtype, g, norm, weight, bias, out_dim):
h_out = conv(g, h, weight=ext_w) h_out = conv(g, h, weight=ext_w)
assert h_out.shape == (ndst, out_dim) assert h_out.shape == (ndst, out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['has_scalar_e_feature'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['has_scalar_e_feature'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right']) @pytest.mark.parametrize('norm', ['none', 'both', 'right'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -119,7 +119,7 @@ def test_graph_conv_e_weight(idtype, g, norm, weight, bias, out_dim): ...@@ -119,7 +119,7 @@ def test_graph_conv_e_weight(idtype, g, norm, weight, bias, out_dim):
h_out = conv(g, h, weight=ext_w, edge_weight=e_w) h_out = conv(g, h, weight=ext_w, edge_weight=e_w)
assert h_out.shape == (ndst, out_dim) assert h_out.shape == (ndst, out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['has_scalar_e_feature'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['has_scalar_e_feature'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right']) @pytest.mark.parametrize('norm', ['none', 'both', 'right'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -144,7 +144,7 @@ def test_graph_conv_e_weight_norm(idtype, g, norm, weight, bias, out_dim): ...@@ -144,7 +144,7 @@ def test_graph_conv_e_weight_norm(idtype, g, norm, weight, bias, out_dim):
h_out = conv(g, h, weight=ext_w, edge_weight=norm_weight) h_out = conv(g, h, weight=ext_w, edge_weight=norm_weight)
assert h_out.shape == (ndst, out_dim) assert h_out.shape == (ndst, out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right']) @pytest.mark.parametrize('norm', ['none', 'both', 'right'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -356,7 +356,7 @@ def test_set_trans(): ...@@ -356,7 +356,7 @@ def test_set_trans():
h2 = st_dec(bg, h1) h2 = st_dec(bg, h1)
assert h2.shape[0] == 3 and h2.shape[1] == 200 and h2.dim() == 2 assert h2.shape[0] == 3 and h2.shape[1] == 200 and h2.dim() == 2
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('O', [1, 8, 32]) @pytest.mark.parametrize('O', [1, 8, 32])
def test_rgcn(idtype, O): def test_rgcn(idtype, O):
ctx = F.ctx() ctx = F.ctx()
...@@ -413,7 +413,7 @@ def test_rgcn(idtype, O): ...@@ -413,7 +413,7 @@ def test_rgcn(idtype, O):
assert h_new.shape == (100, O) assert h_new.shape == (100, O)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 5]) @pytest.mark.parametrize('out_dim', [1, 5])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -437,7 +437,7 @@ def test_gat_conv(g, idtype, out_dim, num_heads): ...@@ -437,7 +437,7 @@ def test_gat_conv(g, idtype, out_dim, num_heads):
gat = gat.to(ctx) gat = gat.to(ctx)
h = gat(g, feat) h = gat(g, feat)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -452,7 +452,7 @@ def test_gat_conv_bi(g, idtype, out_dim, num_heads): ...@@ -452,7 +452,7 @@ def test_gat_conv_bi(g, idtype, out_dim, num_heads):
_, a = gat(g, feat, get_attention=True) _, a = gat(g, feat, get_attention=True)
assert a.shape == (g.number_of_edges(), num_heads, 1) assert a.shape == (g.number_of_edges(), num_heads, 1)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 5]) @pytest.mark.parametrize('out_dim', [1, 5])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -476,7 +476,7 @@ def test_gatv2_conv(g, idtype, out_dim, num_heads): ...@@ -476,7 +476,7 @@ def test_gatv2_conv(g, idtype, out_dim, num_heads):
gat = gat.to(ctx) gat = gat.to(ctx)
h = gat(g, feat) h = gat(g, feat)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -491,7 +491,7 @@ def test_gatv2_conv_bi(g, idtype, out_dim, num_heads): ...@@ -491,7 +491,7 @@ def test_gatv2_conv_bi(g, idtype, out_dim, num_heads):
_, a = gat(g, feat, get_attention=True) _, a = gat(g, feat, get_attention=True)
assert a.shape == (g.number_of_edges(), num_heads, 1) assert a.shape == (g.number_of_edges(), num_heads, 1)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_node_feats', [1, 5]) @pytest.mark.parametrize('out_node_feats', [1, 5])
@pytest.mark.parametrize('out_edge_feats', [1, 5]) @pytest.mark.parametrize('out_edge_feats', [1, 5])
...@@ -513,7 +513,7 @@ def test_egat_conv(g, idtype, out_node_feats, out_edge_feats, num_heads): ...@@ -513,7 +513,7 @@ def test_egat_conv(g, idtype, out_node_feats, out_edge_feats, num_heads):
th.save(egat, tmp_buffer) th.save(egat, tmp_buffer)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite']))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn', 'lstm']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn', 'lstm'])
def test_sage_conv(idtype, g, aggre_type): def test_sage_conv(idtype, g, aggre_type):
...@@ -526,7 +526,7 @@ def test_sage_conv(idtype, g, aggre_type): ...@@ -526,7 +526,7 @@ def test_sage_conv(idtype, g, aggre_type):
h = sage(g, feat) h = sage(g, feat)
assert h.shape[-1] == 10 assert h.shape[-1] == 10
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize('g', get_cases(['bipartite']))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn', 'lstm']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn', 'lstm'])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
...@@ -540,7 +540,7 @@ def test_sage_conv_bi(idtype, g, aggre_type, out_dim): ...@@ -540,7 +540,7 @@ def test_sage_conv_bi(idtype, g, aggre_type, out_dim):
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == g.number_of_dst_nodes() assert h.shape[0] == g.number_of_dst_nodes()
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_sage_conv2(idtype, out_dim): def test_sage_conv2(idtype, out_dim):
# TODO: add test for blocks # TODO: add test for blocks
...@@ -562,7 +562,7 @@ def test_sage_conv2(idtype, out_dim): ...@@ -562,7 +562,7 @@ def test_sage_conv2(idtype, out_dim):
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == 3 assert h.shape[0] == 3
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_sgc_conv(g, idtype, out_dim): def test_sgc_conv(g, idtype, out_dim):
...@@ -588,7 +588,7 @@ def test_sgc_conv(g, idtype, out_dim): ...@@ -588,7 +588,7 @@ def test_sgc_conv(g, idtype, out_dim):
assert F.allclose(h_0, h_1) assert F.allclose(h_0, h_1)
assert h_0.shape[-1] == out_dim assert h_0.shape[-1] == out_dim
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_appnp_conv(g, idtype): def test_appnp_conv(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -604,7 +604,7 @@ def test_appnp_conv(g, idtype): ...@@ -604,7 +604,7 @@ def test_appnp_conv(g, idtype):
assert h.shape[-1] == 5 assert h.shape[-1] == 5
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_appnp_conv_e_weight(g, idtype): def test_appnp_conv_e_weight(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -617,7 +617,7 @@ def test_appnp_conv_e_weight(g, idtype): ...@@ -617,7 +617,7 @@ def test_appnp_conv_e_weight(g, idtype):
h = appnp(g, feat, edge_weight=eweight) h = appnp(g, feat, edge_weight=eweight)
assert h.shape[-1] == 5 assert h.shape[-1] == 5
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_gcn2conv_e_weight(g, idtype): def test_gcn2conv_e_weight(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -632,7 +632,7 @@ def test_gcn2conv_e_weight(g, idtype): ...@@ -632,7 +632,7 @@ def test_gcn2conv_e_weight(g, idtype):
assert h.shape[-1] == 5 assert h.shape[-1] == 5
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_sgconv_e_weight(g, idtype): def test_sgconv_e_weight(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -644,7 +644,7 @@ def test_sgconv_e_weight(g, idtype): ...@@ -644,7 +644,7 @@ def test_sgconv_e_weight(g, idtype):
h = sgconv(g, feat, edge_weight=eweight) h = sgconv(g, feat, edge_weight=eweight)
assert h.shape[-1] == 5 assert h.shape[-1] == 5
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_tagconv_e_weight(g, idtype): def test_tagconv_e_weight(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -657,7 +657,7 @@ def test_tagconv_e_weight(g, idtype): ...@@ -657,7 +657,7 @@ def test_tagconv_e_weight(g, idtype):
h = conv(g, feat, edge_weight=eweight) h = conv(g, feat, edge_weight=eweight)
assert h.shape[-1] == 5 assert h.shape[-1] == 5
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum'])
def test_gin_conv(g, idtype, aggregator_type): def test_gin_conv(g, idtype, aggregator_type):
...@@ -682,7 +682,7 @@ def test_gin_conv(g, idtype, aggregator_type): ...@@ -682,7 +682,7 @@ def test_gin_conv(g, idtype, aggregator_type):
gin = gin.to(ctx) gin = gin.to(ctx)
h = gin(g, feat) h = gin(g, feat)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite']))
def test_gine_conv(g, idtype): def test_gine_conv(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -705,7 +705,7 @@ def test_gine_conv(g, idtype): ...@@ -705,7 +705,7 @@ def test_gine_conv(g, idtype):
gine = gine.to(ctx) gine = gine.to(ctx)
h = gine(g, nfeat, efeat) h = gine(g, nfeat, efeat)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum'])
def test_gin_conv_bi(g, idtype, aggregator_type): def test_gin_conv_bi(g, idtype, aggregator_type):
...@@ -720,7 +720,7 @@ def test_gin_conv_bi(g, idtype, aggregator_type): ...@@ -720,7 +720,7 @@ def test_gin_conv_bi(g, idtype, aggregator_type):
h = gin(g, feat) h = gin(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 12) assert h.shape == (g.number_of_dst_nodes(), 12)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
def test_agnn_conv(g, idtype): def test_agnn_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -731,7 +731,7 @@ def test_agnn_conv(g, idtype): ...@@ -731,7 +731,7 @@ def test_agnn_conv(g, idtype):
h = agnn(g, feat) h = agnn(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 5) assert h.shape == (g.number_of_dst_nodes(), 5)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
def test_agnn_conv_bi(g, idtype): def test_agnn_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -742,7 +742,7 @@ def test_agnn_conv_bi(g, idtype): ...@@ -742,7 +742,7 @@ def test_agnn_conv_bi(g, idtype):
h = agnn(g, feat) h = agnn(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 5) assert h.shape == (g.number_of_dst_nodes(), 5)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_gated_graph_conv(g, idtype): def test_gated_graph_conv(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -757,7 +757,7 @@ def test_gated_graph_conv(g, idtype): ...@@ -757,7 +757,7 @@ def test_gated_graph_conv(g, idtype):
# current we only do shape check # current we only do shape check
assert h.shape[-1] == 10 assert h.shape[-1] == 10
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_gated_graph_conv_one_etype(g, idtype): def test_gated_graph_conv_one_etype(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -774,7 +774,7 @@ def test_gated_graph_conv_one_etype(g, idtype): ...@@ -774,7 +774,7 @@ def test_gated_graph_conv_one_etype(g, idtype):
assert F.allclose(h, h2) assert F.allclose(h, h2)
assert h.shape[-1] == 10 assert h.shape[-1] == 10
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
def test_nn_conv(g, idtype): def test_nn_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -788,7 +788,7 @@ def test_nn_conv(g, idtype): ...@@ -788,7 +788,7 @@ def test_nn_conv(g, idtype):
# currently we only do shape check # currently we only do shape check
assert h.shape[-1] == 10 assert h.shape[-1] == 10
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
def test_nn_conv_bi(g, idtype): def test_nn_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -803,7 +803,7 @@ def test_nn_conv_bi(g, idtype): ...@@ -803,7 +803,7 @@ def test_nn_conv_bi(g, idtype):
# currently we only do shape check # currently we only do shape check
assert h.shape[-1] == 10 assert h.shape[-1] == 10
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_gmm_conv(g, idtype): def test_gmm_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -816,7 +816,7 @@ def test_gmm_conv(g, idtype): ...@@ -816,7 +816,7 @@ def test_gmm_conv(g, idtype):
# currently we only do shape check # currently we only do shape check
assert h.shape[-1] == 10 assert h.shape[-1] == 10
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite', 'block-bipartite'], exclude=['zero-degree']))
def test_gmm_conv_bi(g, idtype): def test_gmm_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -830,7 +830,7 @@ def test_gmm_conv_bi(g, idtype): ...@@ -830,7 +830,7 @@ def test_gmm_conv_bi(g, idtype):
# currently we only do shape check # currently we only do shape check
assert h.shape[-1] == 10 assert h.shape[-1] == 10
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('norm_type', ['both', 'right', 'none']) @pytest.mark.parametrize('norm_type', ['both', 'right', 'none'])
@pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
...@@ -850,7 +850,7 @@ def test_dense_graph_conv(norm_type, g, idtype, out_dim): ...@@ -850,7 +850,7 @@ def test_dense_graph_conv(norm_type, g, idtype, out_dim):
out_dense_conv = dense_conv(adj, feat) out_dense_conv = dense_conv(adj, feat)
assert F.allclose(out_conv, out_dense_conv) assert F.allclose(out_conv, out_dense_conv)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'bipartite']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_dense_sage_conv(g, idtype, out_dim): def test_dense_sage_conv(g, idtype, out_dim):
...@@ -874,7 +874,7 @@ def test_dense_sage_conv(g, idtype, out_dim): ...@@ -874,7 +874,7 @@ def test_dense_sage_conv(g, idtype, out_dim):
out_dense_sage = dense_sage(adj, feat) out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage), g assert F.allclose(out_sage, out_dense_sage), g
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_edge_conv(g, idtype, out_dim): def test_edge_conv(g, idtype, out_dim):
...@@ -890,7 +890,7 @@ def test_edge_conv(g, idtype, out_dim): ...@@ -890,7 +890,7 @@ def test_edge_conv(g, idtype, out_dim):
h1 = edge_conv(g, h0) h1 = edge_conv(g, h0)
assert h1.shape == (g.number_of_dst_nodes(), out_dim) assert h1.shape == (g.number_of_dst_nodes(), out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_edge_conv_bi(g, idtype, out_dim): def test_edge_conv_bi(g, idtype, out_dim):
...@@ -903,7 +903,7 @@ def test_edge_conv_bi(g, idtype, out_dim): ...@@ -903,7 +903,7 @@ def test_edge_conv_bi(g, idtype, out_dim):
h1 = edge_conv(g, (h0, x0)) h1 = edge_conv(g, (h0, x0))
assert h1.shape == (g.number_of_dst_nodes(), out_dim) assert h1.shape == (g.number_of_dst_nodes(), out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -922,7 +922,7 @@ def test_dotgat_conv(g, idtype, out_dim, num_heads): ...@@ -922,7 +922,7 @@ def test_dotgat_conv(g, idtype, out_dim, num_heads):
_, a = dotgat(g, feat, get_attention=True) _, a = dotgat(g, feat, get_attention=True)
assert a.shape == (g.number_of_edges(), num_heads, 1) assert a.shape == (g.number_of_edges(), num_heads, 1)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -1008,7 +1008,7 @@ def test_sequential(): ...@@ -1008,7 +1008,7 @@ def test_sequential():
n_feat = net([g1, g2, g3], n_feat) n_feat = net([g1, g2, g3], n_feat)
assert n_feat.shape == (4, 4) assert n_feat.shape == (4, 4)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_atomic_conv(g, idtype): def test_atomic_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
...@@ -1029,7 +1029,7 @@ def test_atomic_conv(g, idtype): ...@@ -1029,7 +1029,7 @@ def test_atomic_conv(g, idtype):
# current we only do shape check # current we only do shape check
assert h.shape[-1] == 4 assert h.shape[-1] == 4
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 3]) @pytest.mark.parametrize('out_dim', [1, 3])
def test_cf_conv(g, idtype, out_dim): def test_cf_conv(g, idtype, out_dim):
...@@ -1061,7 +1061,7 @@ def myagg(alist, dsttype): ...@@ -1061,7 +1061,7 @@ def myagg(alist, dsttype):
rst = rst + (i + 1) * alist[i] rst = rst + (i + 1) * alist[i]
return rst return rst
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg]) @pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg])
def test_hetero_conv(agg, idtype): def test_hetero_conv(agg, idtype):
g = dgl.heterograph({ g = dgl.heterograph({
...@@ -1186,7 +1186,7 @@ def test_hetero_embedding(out_dim): ...@@ -1186,7 +1186,7 @@ def test_hetero_embedding(out_dim):
assert embeds['user'].shape == (1, out_dim) assert embeds['user'].shape == (1, out_dim)
assert embeds[('user', 'follows', 'user')].shape == (2, out_dim) assert embeds[('user', 'follows', 'user')].shape == (2, out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_gnnexplainer(g, idtype, out_dim): def test_gnnexplainer(g, idtype, out_dim):
...@@ -1317,7 +1317,7 @@ def test_typed_linear(feat_size, regularizer, num_bases): ...@@ -1317,7 +1317,7 @@ def test_typed_linear(feat_size, regularizer, num_bases):
assert th.allclose(y, y_sorted[rev_idx], atol=1e-4, rtol=1e-4) assert th.allclose(y, y_sorted[rev_idx], atol=1e-4, rtol=1e-4)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('in_size', [4]) @pytest.mark.parametrize('in_size', [4])
@pytest.mark.parametrize('num_heads', [1]) @pytest.mark.parametrize('num_heads', [1])
def test_hgt(idtype, in_size, num_heads): def test_hgt(idtype, in_size, num_heads):
...@@ -1433,7 +1433,7 @@ def test_radius_graph(self_loop, get_distances): ...@@ -1433,7 +1433,7 @@ def test_radius_graph(self_loop, get_distances):
if get_distances: if get_distances:
assert th.allclose(dists, dists_target, rtol=1e-03) assert th.allclose(dists, dists_target, rtol=1e-03)
@parametrize_dtype @parametrize_idtype
def test_group_rev_res(idtype): def test_group_rev_res(idtype):
dev = F.ctx() dev = F.ctx()
......
tests/tensorflow/test_nn.py
View file @ 14251504
...@@ -7,7 +7,7 @@ import dgl.nn.tensorflow as nn ...@@ -7,7 +7,7 @@ import dgl.nn.tensorflow as nn
import dgl.function as fn import dgl.function as fn
import backend as F import backend as F
from test_utils.graph_cases import get_cases, random_graph, random_bipartite, random_dglgraph from test_utils.graph_cases import get_cases, random_graph, random_bipartite, random_dglgraph
from test_utils import parametrize_dtype from test_utils import parametrize_idtype
from copy import deepcopy from copy import deepcopy
import numpy as np import numpy as np
...@@ -72,7 +72,7 @@ def test_graph_conv(out_dim): ...@@ -72,7 +72,7 @@ def test_graph_conv(out_dim):
# new_weight = conv.weight.data # new_weight = conv.weight.data
# assert not F.allclose(old_weight, new_weight) # assert not F.allclose(old_weight, new_weight)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right', 'left']) @pytest.mark.parametrize('norm', ['none', 'both', 'right', 'left'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -92,7 +92,7 @@ def test_graph_conv2(idtype, g, norm, weight, bias, out_dim): ...@@ -92,7 +92,7 @@ def test_graph_conv2(idtype, g, norm, weight, bias, out_dim):
h_out = conv(g, h, weight=ext_w) h_out = conv(g, h, weight=ext_w)
assert h_out.shape == (ndst, out_dim) assert h_out.shape == (ndst, out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree', 'dglgraph']))
@pytest.mark.parametrize('norm', ['none', 'both', 'right']) @pytest.mark.parametrize('norm', ['none', 'both', 'right'])
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize('weight', [True, False])
...@@ -262,7 +262,7 @@ def test_rgcn(O): ...@@ -262,7 +262,7 @@ def test_rgcn(O):
assert list(h_new_low.shape) == [100, O] assert list(h_new_low.shape) == [100, O]
assert F.allclose(h_new, h_new_low) assert F.allclose(h_new, h_new_low)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -280,7 +280,7 @@ def test_gat_conv(g, idtype, out_dim, num_heads): ...@@ -280,7 +280,7 @@ def test_gat_conv(g, idtype, out_dim, num_heads):
gat = nn.GATConv(5, out_dim, num_heads, residual=True) gat = nn.GATConv(5, out_dim, num_heads, residual=True)
h = gat(g, feat) h = gat(g, feat)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize('num_heads', [1, 4])
...@@ -294,7 +294,7 @@ def test_gat_conv_bi(g, idtype, out_dim, num_heads): ...@@ -294,7 +294,7 @@ def test_gat_conv_bi(g, idtype, out_dim, num_heads):
_, a = gat(g, feat, get_attention=True) _, a = gat(g, feat, get_attention=True)
assert a.shape == (g.number_of_edges(), num_heads, 1) assert a.shape == (g.number_of_edges(), num_heads, 1)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite']))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn'])
@pytest.mark.parametrize('out_dim', [1, 10]) @pytest.mark.parametrize('out_dim', [1, 10])
...@@ -305,7 +305,7 @@ def test_sage_conv(idtype, g, aggre_type, out_dim): ...@@ -305,7 +305,7 @@ def test_sage_conv(idtype, g, aggre_type, out_dim):
h = sage(g, feat) h = sage(g, feat)
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize('g', get_cases(['bipartite']))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn'])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
...@@ -318,7 +318,7 @@ def test_sage_conv_bi(idtype, g, aggre_type, out_dim): ...@@ -318,7 +318,7 @@ def test_sage_conv_bi(idtype, g, aggre_type, out_dim):
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == g.number_of_dst_nodes() assert h.shape[0] == g.number_of_dst_nodes()
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn'])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_sage_conv_bi_empty(idtype, aggre_type, out_dim): def test_sage_conv_bi_empty(idtype, aggre_type, out_dim):
...@@ -337,7 +337,7 @@ def test_sage_conv_bi_empty(idtype, aggre_type, out_dim): ...@@ -337,7 +337,7 @@ def test_sage_conv_bi_empty(idtype, aggre_type, out_dim):
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == 3 assert h.shape[0] == 3
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_sgc_conv(g, idtype, out_dim): def test_sgc_conv(g, idtype, out_dim):
...@@ -357,7 +357,7 @@ def test_sgc_conv(g, idtype, out_dim): ...@@ -357,7 +357,7 @@ def test_sgc_conv(g, idtype, out_dim):
assert F.allclose(h_0, h_1) assert F.allclose(h_0, h_1)
assert h_0.shape[-1] == out_dim assert h_0.shape[-1] == out_dim
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_appnp_conv(g, idtype): def test_appnp_conv(g, idtype):
ctx = F.ctx() ctx = F.ctx()
...@@ -368,7 +368,7 @@ def test_appnp_conv(g, idtype): ...@@ -368,7 +368,7 @@ def test_appnp_conv(g, idtype):
h = appnp(g, feat) h = appnp(g, feat)
assert h.shape[-1] == 5 assert h.shape[-1] == 5
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite']))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum'])
def test_gin_conv(g, idtype, aggregator_type): def test_gin_conv(g, idtype, aggregator_type):
...@@ -382,7 +382,7 @@ def test_gin_conv(g, idtype, aggregator_type): ...@@ -382,7 +382,7 @@ def test_gin_conv(g, idtype, aggregator_type):
h = gin(g, feat) h = gin(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 12) assert h.shape == (g.number_of_dst_nodes(), 12)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize('g', get_cases(['bipartite']))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum'])
def test_gin_conv_bi(g, idtype, aggregator_type): def test_gin_conv_bi(g, idtype, aggregator_type):
...@@ -395,7 +395,7 @@ def test_gin_conv_bi(g, idtype, aggregator_type): ...@@ -395,7 +395,7 @@ def test_gin_conv_bi(g, idtype, aggregator_type):
h = gin(g, feat) h = gin(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 12) assert h.shape == (g.number_of_dst_nodes(), 12)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_edge_conv(g, idtype, out_dim): def test_edge_conv(g, idtype, out_dim):
...@@ -406,7 +406,7 @@ def test_edge_conv(g, idtype, out_dim): ...@@ -406,7 +406,7 @@ def test_edge_conv(g, idtype, out_dim):
h1 = edge_conv(g, h0) h1 = edge_conv(g, h0)
assert h1.shape == (g.number_of_dst_nodes(), out_dim) assert h1.shape == (g.number_of_dst_nodes(), out_dim)
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree']))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize('out_dim', [1, 2])
def test_edge_conv_bi(g, idtype, out_dim): def test_edge_conv_bi(g, idtype, out_dim):
...@@ -425,7 +425,7 @@ def myagg(alist, dsttype): ...@@ -425,7 +425,7 @@ def myagg(alist, dsttype):
rst = rst + (i + 1) * alist[i] rst = rst + (i + 1) * alist[i]
return rst return rst
@parametrize_dtype @parametrize_idtype
@pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg]) @pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg])
def test_hetero_conv(agg, idtype): def test_hetero_conv(agg, idtype):
g = dgl.heterograph({ g = dgl.heterograph({
......
tests/test_utils/__init__.py
View file @ 14251504
import pytest import pytest
import backend as F import backend as F
if F._default_context_str == 'cpu': parametrize_idtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
parametrize_dtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
else:
# only test int32 on GPU because many graph operators are not supported for int64.
parametrize_dtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
from .checks import * from .checks import *
from .graph_cases import get_cases from .graph_cases import get_cases
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