def test_graph_conv(): in_channels, out_channels = (16, 32) edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]]) edge_weight = torch.randn(edge_index.size(1)) num_nodes = edge_index.max().item() + 1 x = torch.randn((num_nodes, in_channels)) conv = GraphConv(in_channels, out_channels) assert conv.__repr__() == 'GraphConv(16, 32)' out1 = conv(x, edge_index) assert out1.size() == (num_nodes, out_channels) out2 = conv(x, edge_index, edge_weight) assert out2.size() == (num_nodes, out_channels) jit_conv = conv.jittable(x=x, edge_index=edge_index) jit_conv = torch.jit.script(jit_conv) assert jit_conv(x, edge_index).tolist() == out1.tolist() assert jit_conv(x, edge_index, edge_weight).tolist() == out2.tolist()
def test_graph_conv(): x1 = torch.randn(4, 8) x2 = torch.randn(2, 16) edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]]) row, col = edge_index value = torch.randn(edge_index.size(1)) adj1 = SparseTensor(row=row, col=col, sparse_sizes=(4, 4)) adj2 = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4)) conv = GraphConv(8, 32) assert conv.__repr__() == 'GraphConv(8, 32)' out11 = conv(x1, edge_index) assert out11.size() == (4, 32) assert conv(x1, edge_index, size=(4, 4)).tolist() == out11.tolist() assert conv(x1, adj1.t()).tolist() == out11.tolist() out12 = conv(x1, edge_index, value) assert out12.size() == (4, 32) assert conv(x1, edge_index, value, size=(4, 4)).tolist() == out12.tolist() assert conv(x1, adj2.t()).tolist() == out12.tolist() t = '(Tensor, Tensor, OptTensor, Size) -> Tensor' jit = torch.jit.script(conv.jittable(t)) assert jit(x1, edge_index).tolist() == out11.tolist() assert jit(x1, edge_index, size=(4, 4)).tolist() == out11.tolist() assert jit(x1, edge_index, value).tolist() == out12.tolist() assert jit(x1, edge_index, value, size=(4, 4)).tolist() == out12.tolist() t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor' jit = torch.jit.script(conv.jittable(t)) assert jit(x1, adj1.t()).tolist() == out11.tolist() assert jit(x1, adj2.t()).tolist() == out12.tolist() adj1 = adj1.sparse_resize((4, 2)) adj2 = adj2.sparse_resize((4, 2)) conv = GraphConv((8, 16), 32) assert conv.__repr__() == 'GraphConv((8, 16), 32)' out21 = conv((x1, x2), edge_index) out22 = conv((x1, x2), edge_index, value) out23 = conv((x1, None), edge_index, size=(4, 2)) out24 = conv((x1, None), edge_index, value, size=(4, 2)) assert out21.size() == (2, 32) assert out22.size() == (2, 32) assert out23.size() == (2, 32) assert out24.size() == (2, 32) assert conv((x1, x2), edge_index, size=(4, 2)).tolist() == out21.tolist() assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out22.tolist() assert conv((x1, x2), adj1.t()).tolist() == out21.tolist() assert conv((x1, x2), adj2.t()).tolist() == out22.tolist() t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor' jit = torch.jit.script(conv.jittable(t)) assert jit((x1, x2), edge_index).tolist() == out21.tolist() assert jit((x1, x2), edge_index, size=(4, 2)).tolist() == out21.tolist() assert jit((x1, x2), edge_index, value).tolist() == out22.tolist() assert jit((x1, x2), edge_index, value, (4, 2)).tolist() == out22.tolist() assert jit((x1, None), edge_index, size=(4, 2)).tolist() == out23.tolist() assert jit((x1, None), edge_index, value, (4, 2)).tolist() == out24.tolist() t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor' jit = torch.jit.script(conv.jittable(t)) assert jit((x1, x2), adj1.t()).tolist() == out21.tolist() assert jit((x1, x2), adj2.t()).tolist() == out22.tolist() assert jit((x1, None), adj1.t()).tolist() == out23.tolist() assert jit((x1, None), adj2.t()).tolist() == out24.tolist()