def test_gen_conv(aggr):
x1 = torch.randn(4, 16)
x2 = torch.randn(2, 16)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
value = torch.randn(row.size(0), 16)
adj1 = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
adj2 = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = GENConv(16, 32, aggr)
assert conv.__repr__() == f'GENConv(16, 32, aggr={aggr})'
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, (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))
out21 = conv((x1, x2), edge_index)
assert out21.size() == (2, 32)
assert conv((x1, x2), edge_index, size=(4, 2)).tolist() == out21.tolist()
assert conv((x1, x2), adj1.t()).tolist() == out21.tolist()
out22 = conv((x1, x2), edge_index, value)
assert out22.size() == (2, 32)
assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out22.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()
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()
def test_feast_conv():
x1 = torch.randn(4, 16)
x2 = torch.randn(2, 16)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = FeaStConv(16, 32, heads=2)
assert conv.__repr__() == 'FeaStConv(16, 32, heads=2)'
out = conv(x1, edge_index)
assert out.size() == (4, 32)
assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
t = '(Tensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
t = '(Tensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
adj = adj.sparse_resize((4, 2))
out = conv((x1, x2), edge_index)
assert out.size() == (2, 32)
assert torch.allclose(conv((x1, x2), adj.t()), out, atol=1e-6)
t = '(PairTensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out.tolist()
t = '(PairTensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), adj.t()), out, atol=1e-6)
def test_res_gated_graph_conv():
x1 = torch.randn(4, 8)
x2 = torch.randn(2, 32)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = ResGatedGraphConv(8, 32)
assert conv.__repr__() == 'ResGatedGraphConv(8, 32)'
out = conv(x1, edge_index)
assert out.size() == (4, 32)
assert conv(x1, adj.t()).tolist() == out.tolist()
t = '(Tensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
t = '(Tensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
conv = ResGatedGraphConv((8, 32), 32)
assert conv.__repr__() == 'ResGatedGraphConv((8, 32), 32)'
out = conv((x1, x2), edge_index)
assert out.size() == (2, 32)
assert conv((x1, x2), adj.t()).tolist() == out.tolist()
t = '(PairTensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out.tolist()
t = '(PairTensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out.tolist()
def test_nn_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.rand(row.size(0), 3)
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
nn = Seq(Lin(3, 32), ReLU(), Lin(32, 8 * 32))
conv = NNConv(8, 32, nn=nn)
assert conv.__repr__() == (
'NNConv(8, 32, aggr=add, nn=Sequential(\n'
' (0): Linear(in_features=3, out_features=32, bias=True)\n'
' (1): ReLU()\n'
' (2): Linear(in_features=32, out_features=256, bias=True)\n'
'))')
out = conv(x1, edge_index, value)
assert out.size() == (4, 32)
assert conv(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
if is_full_test():
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index, value).tolist() == out.tolist()
assert jit(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
conv = NNConv((8, 16), 32, nn=nn)
assert conv.__repr__() == (
'NNConv((8, 16), 32, aggr=add, nn=Sequential(\n'
' (0): Linear(in_features=3, out_features=32, bias=True)\n'
' (1): ReLU()\n'
' (2): Linear(in_features=32, out_features=256, bias=True)\n'
'))')
out1 = conv((x1, x2), edge_index, value)
out2 = conv((x1, None), edge_index, value, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out1.tolist()
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
if is_full_test():
t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index, value).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, value,
size=(4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index, value,
size=(4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
def test_my_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(row.size(0))
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = MyConv(8, 32)
out = conv(x1, edge_index, value)
assert out.size() == (4, 32)
assert conv(x1, edge_index, value, (4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
conv.fuse = False
assert conv(x1, adj.t()).tolist() == out.tolist()
conv.fuse = True
adj = adj.sparse_resize((4, 2))
conv = MyConv((8, 16), 32)
out1 = conv((x1, x2), edge_index, value)
out2 = conv((x1, None), edge_index, value, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out1.tolist()
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
conv.fuse = False
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
conv.fuse = True
def test_signed_conv():
x = torch.randn(4, 16)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv1 = SignedConv(16, 32, first_aggr=True)
assert conv1.__repr__() == 'SignedConv(16, 32, first_aggr=True)'
conv2 = SignedConv(32, 48, first_aggr=False)
assert conv2.__repr__() == 'SignedConv(32, 48, first_aggr=False)'
out1 = conv1(x, edge_index, edge_index)
assert out1.size() == (4, 64)
assert conv1(x, adj.t(), adj.t()).tolist() == out1.tolist()
out2 = conv2(out1, edge_index, edge_index)
assert out2.size() == (4, 96)
assert conv2(out1, adj.t(), adj.t()).tolist() == out2.tolist()
if is_full_test():
t = '(Tensor, Tensor, Tensor) -> Tensor'
jit1 = torch.jit.script(conv1.jittable(t))
jit2 = torch.jit.script(conv2.jittable(t))
assert jit1(x, edge_index, edge_index).tolist() == out1.tolist()
assert jit2(out1, edge_index, edge_index).tolist() == out2.tolist()
t = '(Tensor, SparseTensor, SparseTensor) -> Tensor'
jit1 = torch.jit.script(conv1.jittable(t))
jit2 = torch.jit.script(conv2.jittable(t))
assert jit1(x, adj.t(), adj.t()).tolist() == out1.tolist()
assert jit2(out1, adj.t(), adj.t()).tolist() == out2.tolist()
adj = adj.sparse_resize((4, 2))
assert torch.allclose(conv1((x, x[:2]), edge_index, edge_index), out1[:2],
atol=1e-6)
assert torch.allclose(conv1((x, x[:2]), adj.t(), adj.t()), out1[:2],
atol=1e-6)
assert torch.allclose(conv2((out1, out1[:2]), edge_index, edge_index),
out2[:2], atol=1e-6)
assert torch.allclose(conv2((out1, out1[:2]), adj.t(), adj.t()), out2[:2],
atol=1e-6)
if is_full_test():
t = '(PairTensor, Tensor, Tensor) -> Tensor'
jit1 = torch.jit.script(conv1.jittable(t))
jit2 = torch.jit.script(conv2.jittable(t))
assert torch.allclose(jit1((x, x[:2]), edge_index, edge_index),
out1[:2], atol=1e-6)
assert torch.allclose(jit2((out1, out1[:2]), edge_index, edge_index),
out2[:2], atol=1e-6)
t = '(PairTensor, SparseTensor, SparseTensor) -> Tensor'
jit1 = torch.jit.script(conv1.jittable(t))
jit2 = torch.jit.script(conv2.jittable(t))
assert torch.allclose(jit1((x, x[:2]), adj.t(), adj.t()), out1[:2],
atol=1e-6)
assert torch.allclose(jit2((out1, out1[:2]), adj.t(), adj.t()),
out2[:2], atol=1e-6)
def test_edge_conv_conv():
x1 = torch.randn(4, 16)
x2 = torch.randn(2, 16)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
nn = Seq(Lin(32, 16), ReLU(), Lin(16, 32))
conv = EdgeConv(nn)
assert conv.__repr__() == (
'EdgeConv(nn=Sequential(\n'
' (0): Linear(in_features=32, out_features=16, bias=True)\n'
' (1): ReLU()\n'
' (2): Linear(in_features=16, out_features=32, bias=True)\n'
'))')
out1 = conv(x1, edge_index)
assert out1.size() == (4, 32)
assert conv((x1, x1), edge_index).tolist() == out1.tolist()
assert conv(x1, adj.t()).tolist() == out1.tolist()
assert conv((x1, x1), adj.t()).tolist() == out1.tolist()
adj = adj.sparse_resize((4, 2))
out2 = conv((x1, x2), edge_index)
assert out2.size() == (2, 32)
assert conv((x1, x2), adj.t()).tolist() == out2.tolist()
if is_full_test():
t = '(Tensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out1.tolist()
t = '(PairTensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x1), edge_index).tolist() == out1.tolist()
assert jit((x1, x2), edge_index).tolist() == out2.tolist()
adj = adj.sparse_resize((4, 4))
t = '(Tensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out1.tolist()
t = '(PairTensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x1), adj.t()).tolist() == out1.tolist()
adj = adj.sparse_resize((4, 2))
assert jit((x1, x2), adj.t()).tolist() == out2.tolist()
def test_gmm_conv(separate_gaussians):
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.rand(row.size(0), 3)
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = GMMConv(8,
32,
dim=3,
kernel_size=25,
separate_gaussians=separate_gaussians)
assert conv.__repr__() == 'GMMConv(8, 32, dim=3)'
out = conv(x1, edge_index, value)
assert out.size() == (4, 32)
assert torch.allclose(conv(x1, edge_index, value, size=(4, 4)), out)
assert torch.allclose(conv(x1, adj.t()), out)
if is_full_test():
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, edge_index, value), out)
assert torch.allclose(jit(x1, edge_index, value, size=(4, 4)), out)
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, adj.t()), out)
adj = adj.sparse_resize((4, 2))
conv = GMMConv((8, 16),
32,
dim=3,
kernel_size=5,
separate_gaussians=separate_gaussians)
assert conv.__repr__() == 'GMMConv((8, 16), 32, dim=3)'
out1 = conv((x1, x2), edge_index, value)
out2 = conv((x1, None), edge_index, value, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert torch.allclose(conv((x1, x2), edge_index, value, (4, 2)), out1)
assert torch.allclose(conv((x1, x2), adj.t()), out1)
assert torch.allclose(conv((x1, None), adj.t()), out2)
if is_full_test():
t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), edge_index, value), out1)
assert torch.allclose(jit((x1, x2), edge_index, value, size=(4, 2)),
out1)
assert torch.allclose(jit((x1, None), edge_index, value, size=(4, 2)),
out2)
t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), adj.t()), out1)
assert torch.allclose(jit((x1, None), adj.t()), out2)
def test_ppf_conv():
x1 = torch.randn(4, 16)
pos1 = torch.randn(4, 3)
pos2 = torch.randn(2, 3)
n1 = F.normalize(torch.rand(4, 3), dim=-1)
n2 = F.normalize(torch.rand(2, 3), dim=-1)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
local_nn = Seq(Lin(16 + 4, 32), ReLU(), Lin(32, 32))
global_nn = Seq(Lin(32, 32))
conv = PPFConv(local_nn, global_nn)
assert conv.__repr__() == (
'PPFConv(local_nn=Sequential(\n'
' (0): Linear(in_features=20, out_features=32, bias=True)\n'
' (1): ReLU()\n'
' (2): Linear(in_features=32, out_features=32, bias=True)\n'
'), global_nn=Sequential(\n'
' (0): Linear(in_features=32, out_features=32, bias=True)\n'
'))')
out = conv(x1, pos1, n1, edge_index)
assert out.size() == (4, 32)
assert torch.allclose(conv(x1, pos1, n1, adj.t()), out, atol=1e-6)
t = '(OptTensor, Tensor, Tensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, pos1, n1, edge_index).tolist() == out.tolist()
t = '(OptTensor, Tensor, Tensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, pos1, n1, adj.t()), out, atol=1e-6)
adj = adj.sparse_resize((4, 2))
out = conv(x1, (pos1, pos2), (n1, n2), edge_index)
assert out.size() == (2, 32)
assert conv((x1, None), (pos1, pos2), (n1, n2),
edge_index).tolist() == out.tolist()
assert torch.allclose(conv(x1, (pos1, pos2), (n1, n2), adj.t()),
out,
atol=1e-6)
assert torch.allclose(conv((x1, None), (pos1, pos2), (n1, n2), adj.t()),
out,
atol=1e-6)
t = '(PairOptTensor, PairTensor, PairTensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, None), (pos1, pos2), (n1, n2),
edge_index).tolist() == out.tolist()
t = '(PairOptTensor, PairTensor, PairTensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, None), (pos1, pos2), (n1, n2), adj.t()),
out,
atol=1e-6)
def test_cg_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
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = CGConv(8)
assert conv.__repr__() == 'CGConv(8, dim=0)'
out = conv(x1, edge_index)
assert out.size() == (4, 8)
assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
assert jit(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
conv = CGConv((8, 16))
assert conv.__repr__() == 'CGConv((8, 16), dim=0)'
out = conv((x1, x2), edge_index)
assert out.size() == (2, 16)
assert conv((x1, x2), edge_index, size=(4, 2)).tolist() == out.tolist()
assert conv((x1, x2), adj.t()).tolist() == out.tolist()
t = '(PairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out.tolist()
assert jit((x1, x2), edge_index, size=(4, 2)).tolist() == out.tolist()
t = '(PairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out.tolist()
# Test batch_norm true:
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = CGConv(8, batch_norm=True)
assert conv.__repr__() == 'CGConv(8, dim=0)'
out = conv(x1, edge_index)
assert out.size() == (4, 8)
assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
assert jit(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
def test_point_transformer_conv():
x1 = torch.rand(4, 16)
x2 = torch.randn(2, 8)
pos1 = torch.rand(4, 3)
pos2 = torch.randn(2, 3)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = PointTransformerConv(in_channels=16, out_channels=32)
assert str(conv) == 'PointTransformerConv(16, 32)'
out = conv(x1, pos1, edge_index)
assert out.size() == (4, 32)
assert torch.allclose(conv(x1, pos1, adj.t()), out, atol=1e-6)
if is_full_test():
t = '(Tensor, Tensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, pos1, edge_index).tolist() == out.tolist()
t = '(Tensor, Tensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, pos1, adj.t()), out, atol=1e-6)
pos_nn = Sequential(Linear(3, 16), ReLU(), Linear(16, 32))
attn_nn = Sequential(Linear(32, 32), ReLU(), Linear(32, 32))
conv = PointTransformerConv(16, 32, pos_nn, attn_nn)
out = conv(x1, pos1, edge_index)
assert out.size() == (4, 32)
assert torch.allclose(conv(x1, pos1, adj.t()), out, atol=1e-6)
conv = PointTransformerConv((16, 8), 32)
adj = adj.sparse_resize((4, 2))
out = conv((x1, x2), (pos1, pos2), edge_index)
assert out.size() == (2, 32)
assert torch.allclose(conv((x1, x2), (pos1, pos2), adj.t()),
out,
atol=1e-6)
if is_full_test():
t = '(PairTensor, PairTensor, Tensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), (pos1, pos2), edge_index).tolist() == out.tolist()
t = '(PairTensor, PairTensor, SparseTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), (pos1, pos2), adj.t()),
out,
atol=1e-6)
def test_gine_conv():
x1 = torch.randn(4, 16)
x2 = torch.randn(2, 16)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
value = torch.randn(row.size(0), 16)
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
nn = Seq(Lin(16, 32), ReLU(), Lin(32, 32))
conv = GINEConv(nn, train_eps=True)
assert conv.__repr__() == (
'GINEConv(nn=Sequential(\n'
' (0): Linear(in_features=16, out_features=32, bias=True)\n'
' (1): ReLU()\n'
' (2): Linear(in_features=32, out_features=32, bias=True)\n'
'))')
out = conv(x1, edge_index, value)
assert out.size() == (4, 32)
assert conv(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index, value).tolist() == out.tolist()
assert jit(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
out1 = conv((x1, x2), edge_index, value)
out2 = conv((x1, None), edge_index, value, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out1.tolist()
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index, value).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, value,
size=(4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index, value,
size=(4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
def test_spline_conv():
warnings.filterwarnings('ignore', '.*non-optimized CPU version.*')
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.rand(row.size(0), 3)
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = SplineConv(8, 32, dim=3, kernel_size=5)
assert conv.__repr__() == 'SplineConv(8, 32, dim=3)'
out = conv(x1, edge_index, value)
assert out.size() == (4, 32)
assert torch.allclose(conv(x1, edge_index, value, size=(4, 4)), out)
assert torch.allclose(conv(x1, adj.t()), out)
if is_full_test():
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, edge_index, value), out)
assert torch.allclose(jit(x1, edge_index, value, size=(4, 4)), out)
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, adj.t()), out)
adj = adj.sparse_resize((4, 2))
conv = SplineConv((8, 16), 32, dim=3, kernel_size=5)
assert conv.__repr__() == 'SplineConv((8, 16), 32, dim=3)'
out1 = conv((x1, x2), edge_index, value)
out2 = conv((x1, None), edge_index, value, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert torch.allclose(conv((x1, x2), edge_index, value, (4, 2)), out1)
assert torch.allclose(conv((x1, x2), adj.t()), out1)
assert torch.allclose(conv((x1, None), adj.t()), out2)
if is_full_test():
t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), edge_index, value), out1)
assert torch.allclose(jit((x1, x2), edge_index, value, size=(4, 2)),
out1)
assert torch.allclose(jit((x1, None), edge_index, value, size=(4, 2)),
out2)
t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), adj.t()), out1)
assert torch.allclose(jit((x1, None), adj.t()), out2)
def test_nn_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.rand(row.size(0), 3)
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
nn = Seq(Lin(3, 32), ReLU(), Lin(32, 8 * 32))
conv = NNConv(8, 32, nn=nn)
assert conv.__repr__() == 'NNConv(8, 32)'
out = conv(x1, edge_index, value)
assert out.size() == (4, 32)
assert conv(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index, value).tolist() == out.tolist()
assert jit(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
conv = NNConv((8, 16), 32, nn=nn)
assert conv.__repr__() == 'NNConv((8, 16), 32)'
out1 = conv((x1, x2), edge_index, value)
out2 = conv((x1, None), edge_index, value, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out1.tolist()
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index, value).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, value,
size=(4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index, value,
size=(4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
def test_sage_conv(project):
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
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = SAGEConv(8, 32, project=project)
assert str(conv) == 'SAGEConv(8, 32, aggr=mean)'
out = conv(x1, edge_index)
assert out.size() == (4, 32)
assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
if is_full_test():
t = '(Tensor, Tensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
assert jit(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
conv = SAGEConv((8, 16), 32, project=project)
assert str(conv) == 'SAGEConv((8, 16), 32, aggr=mean)'
out1 = conv((x1, x2), edge_index)
out2 = conv((x1, None), edge_index, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert conv((x1, x2), edge_index, (4, 2)).tolist() == out1.tolist()
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
if is_full_test():
t = '(OptPairTensor, Tensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, size=(4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index,
size=(4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
def test_sage_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(row.size(0))
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = SAGEConv(8, 32)
assert conv.__repr__() == 'SAGEConv(8, 32)'
out = conv(x1, edge_index)
assert out.size() == (4, 32)
assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
t = '(Tensor, Tensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
assert jit(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
conv = SAGEConv((8, 16), 32)
assert conv.__repr__() == 'SAGEConv((8, 16), 32)'
out1 = conv((x1, x2), edge_index)
out2 = conv((x1, None), edge_index, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert conv((x1, x2), edge_index, (4, 2)).tolist() == out1.tolist()
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
t = '(OptPairTensor, Tensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, size=(4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index, size=(4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
def test_my_conv_jittable():
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(row.size(0))
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = MyConv(8, 32)
out = conv(x1, edge_index, value)
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index, value).tolist() == out.tolist()
assert jit(x1, edge_index, value, (4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
jit.fuse = False
assert jit(x1, adj.t()).tolist() == out.tolist()
jit.fuse = True
adj = adj.sparse_resize((4, 2))
conv = MyConv((8, 16), 32)
out1 = conv((x1, x2), edge_index, value)
out2 = conv((x1, None), edge_index, value, (4, 2))
t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index, value).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, value, (4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index, value, (4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
jit.fuse = False
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
jit.fuse = True
def test_cg_conv_with_edge_features():
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.rand(row.size(0), 3)
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = CGConv(8, dim=3)
assert conv.__repr__() == 'CGConv(8, dim=3)'
out = conv(x1, edge_index, value)
assert out.size() == (4, 8)
assert conv(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index, value).tolist() == out.tolist()
assert jit(x1, edge_index, value, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
conv = CGConv((8, 16), dim=3)
assert conv.__repr__() == 'CGConv((8, 16), dim=3)'
out = conv((x1, x2), edge_index, value)
assert out.size() == (2, 16)
assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out.tolist()
assert conv((x1, x2), adj.t()).tolist() == out.tolist()
t = '(PairTensor, Tensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index, value).tolist() == out.tolist()
assert jit((x1, x2), edge_index, value, (4, 2)).tolist() == out.tolist()
t = '(PairTensor, SparseTensor, OptTensor, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out.tolist()
def test_transformer_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
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = TransformerConv(8, 32, heads=2, beta=True)
assert conv.__repr__() == 'TransformerConv(8, 32, heads=2)'
out = conv(x1, edge_index)
assert out.size() == (4, 64)
assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
t = '(Tensor, Tensor, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
t = '(Tensor, SparseTensor, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
adj = adj.sparse_resize((4, 2))
conv = TransformerConv((8, 16), 32, heads=2, beta=True)
assert conv.__repr__() == 'TransformerConv((8, 16), 32, heads=2)'
out = conv((x1, x2), edge_index)
assert out.size() == (2, 64)
assert torch.allclose(conv((x1, x2), adj.t()), out, atol=1e-6)
t = '(PairTensor, Tensor, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out.tolist()
t = '(PairTensor, SparseTensor, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(conv((x1, x2), adj.t()), out, atol=1e-6)
def test_transformer_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
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = TransformerConv(8, 32, heads=2, beta=True)
assert conv.__repr__() == 'TransformerConv(8, 32, heads=2)'
out = conv(x1, edge_index)
assert out.size() == (4, 64)
assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
if is_full_test():
t = '(Tensor, Tensor, NoneType, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, edge_index), out)
t = '(Tensor, SparseTensor, NoneType, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
# Test `return_attention_weights`.
result = conv(x1, edge_index, return_attention_weights=True)
assert torch.allclose(result[0], out)
assert result[1][0].size() == (2, 4)
assert result[1][1].size() == (4, 2)
assert result[1][1].min() >= 0 and result[1][1].max() <= 1
assert conv._alpha is None
result = conv(x1, adj.t(), return_attention_weights=True)
assert torch.allclose(result[0], out, atol=1e-6)
assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 4
assert conv._alpha is None
if is_full_test():
t = ('(Tensor, Tensor, NoneType, bool) -> '
'Tuple[Tensor, Tuple[Tensor, Tensor]]')
jit = torch.jit.script(conv.jittable(t))
result = jit(x1, edge_index, return_attention_weights=True)
assert torch.allclose(result[0], out)
assert result[1][0].size() == (2, 4)
assert result[1][1].size() == (4, 2)
assert result[1][1].min() >= 0 and result[1][1].max() <= 1
assert conv._alpha is None
t = ('(Tensor, SparseTensor, NoneType, bool) -> '
'Tuple[Tensor, SparseTensor]')
jit = torch.jit.script(conv.jittable(t))
result = jit(x1, adj.t(), return_attention_weights=True)
assert torch.allclose(result[0], out, atol=1e-6)
assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 4
assert conv._alpha is None
adj = adj.sparse_resize((4, 2))
conv = TransformerConv((8, 16), 32, heads=2, beta=True)
assert conv.__repr__() == 'TransformerConv((8, 16), 32, heads=2)'
out = conv((x1, x2), edge_index)
assert out.size() == (2, 64)
assert torch.allclose(conv((x1, x2), adj.t()), out, atol=1e-6)
if is_full_test():
t = '(PairTensor, Tensor, NoneType, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), edge_index), out)
t = '(PairTensor, SparseTensor, NoneType, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), adj.t()), out, atol=1e-6)
def test_film_conv():
x1 = torch.randn(4, 4)
x2 = torch.randn(2, 16)
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [0, 0, 1, 0, 1, 1]])
edge_type = torch.tensor([0, 1, 1, 0, 0, 1])
row, col = edge_index
adj = SparseTensor(row=row, col=col, value=edge_type, sparse_sizes=(4, 4))
conv = FiLMConv(4, 32)
assert conv.__repr__() == 'FiLMConv(4, 32, num_relations=1)'
out1 = conv(x1, edge_index)
assert out1.size() == (4, 32)
assert conv(x1, adj.t().set_value(None)).tolist() == out1.tolist()
t = '(Tensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out1.tolist()
t = '(Tensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert conv(x1, adj.t().set_value(None)).tolist() == out1.tolist()
conv = FiLMConv(4, 32, num_relations=2)
assert conv.__repr__() == 'FiLMConv(4, 32, num_relations=2)'
out1 = conv(x1, edge_index, edge_type)
assert out1.size() == (4, 32)
assert conv(x1, adj.t()).tolist() == out1.tolist()
t = '(Tensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index, edge_type).tolist() == out1.tolist()
t = '(Tensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert conv(x1, adj.t()).tolist() == out1.tolist()
adj = adj.sparse_resize((4, 2))
conv = FiLMConv((4, 16), 32)
assert conv.__repr__() == 'FiLMConv((4, 16), 32, num_relations=1)'
out1 = conv((x1, x2), edge_index)
assert out1.size() == (2, 32)
assert conv((x1, x2), adj.t().set_value(None)).tolist() == out1.tolist()
t = '(PairTensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out1.tolist()
t = '(PairTensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert conv((x1, x2), adj.t().set_value(None)).tolist() == out1.tolist()
conv = FiLMConv((4, 16), 32, num_relations=2)
assert conv.__repr__() == 'FiLMConv((4, 16), 32, num_relations=2)'
out1 = conv((x1, x2), edge_index, edge_type)
assert out1.size() == (2, 32)
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
t = '(PairTensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index, edge_type).tolist() == out1.tolist()
t = '(PairTensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
def test_gat_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
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = GATConv(8, 32, heads=2)
assert conv.__repr__() == 'GATConv(8, 32, heads=2)'
out = conv(x1, edge_index)
assert out.size() == (4, 64)
assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
t = '(Tensor, Tensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
assert jit(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
# Test `return_attention_weights`.
result = conv(x1, edge_index, return_attention_weights=True)
assert result[0].tolist() == out.tolist()
assert result[1][0].size() == (2, 7)
assert result[1][1].size() == (7, 2)
assert result[1][1].min() >= 0 and result[1][1].max() <= 1
assert conv._alpha is None
result = conv(x1, adj.t(), return_attention_weights=True)
assert torch.allclose(result[0], out, atol=1e-6)
assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 7
assert conv._alpha is None
t = ('(Tensor, Tensor, OptTensor, Size, bool) -> '
'Tuple[Tensor, Tuple[Tensor, Tensor]]')
jit = torch.jit.script(conv.jittable(t))
result = jit(x1, edge_index, return_attention_weights=True)
assert result[0].tolist() == out.tolist()
assert result[1][0].size() == (2, 7)
assert result[1][1].size() == (7, 2)
assert result[1][1].min() >= 0 and result[1][1].max() <= 1
assert conv._alpha is None
t = ('(Tensor, SparseTensor, OptTensor, Size, bool) -> '
'Tuple[Tensor, SparseTensor]')
jit = torch.jit.script(conv.jittable(t))
result = jit(x1, adj.t(), return_attention_weights=True)
assert torch.allclose(result[0], out, atol=1e-6)
assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 7
assert conv._alpha is None
adj = adj.sparse_resize((4, 2))
conv = GATConv((8, 16), 32, heads=2)
assert conv.__repr__() == 'GATConv((8, 16), 32, heads=2)'
out1 = conv((x1, x2), edge_index)
out2 = conv((x1, None), edge_index, size=(4, 2))
assert out1.size() == (2, 64)
assert out2.size() == (2, 64)
assert conv((x1, x2), edge_index, size=(4, 2)).tolist() == out1.tolist()
assert torch.allclose(conv((x1, x2), adj.t()), out1, atol=1e-6)
assert torch.allclose(conv((x1, None), adj.t()), out2, atol=1e-6)
t = '(OptPairTensor, Tensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, size=(4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index, size=(4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), adj.t()), out1, atol=1e-6)
assert torch.allclose(jit((x1, None), adj.t()), out2, atol=1e-6)
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()
def test_rgcn_conv(cls, conf):
num_bases, num_blocks = conf
x1 = torch.randn(4, 4)
x2 = torch.randn(2, 16)
idx1 = torch.arange(4)
idx2 = torch.arange(2)
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [0, 0, 1, 0, 1, 1]])
edge_type = torch.tensor([0, 1, 1, 0, 0, 1])
row, col = edge_index
adj = SparseTensor(row=row, col=col, value=edge_type, sparse_sizes=(4, 4))
conv = cls(4, 32, 2, num_bases, num_blocks)
assert conv.__repr__() == f'{cls.__name__}(4, 32, num_relations=2)'
out1 = conv(x1, edge_index, edge_type)
assert out1.size() == (4, 32)
assert conv(x1, adj.t()).tolist() == out1.tolist()
if num_blocks is None:
out2 = conv(None, edge_index, edge_type)
assert out2.size() == (4, 32)
assert conv(None, adj.t()).tolist() == out2.tolist()
t = '(OptTensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index, edge_type).tolist() == out1.tolist()
if num_blocks is None:
assert jit(idx1, edge_index, edge_type).tolist() == out2.tolist()
assert jit(None, edge_index, edge_type).tolist() == out2.tolist()
t = '(OptTensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out1.tolist()
if num_blocks is None:
assert jit(idx1, adj.t()).tolist() == out2.tolist()
assert jit(None, adj.t()).tolist() == out2.tolist()
adj = adj.sparse_resize((4, 2))
conv = cls((4, 16), 32, 2, num_bases, num_blocks)
assert conv.__repr__() == f'{cls.__name__}((4, 16), 32, num_relations=2)'
out1 = conv((x1, x2), edge_index, edge_type)
assert out1.size() == (2, 32)
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
if num_blocks is None:
out2 = conv((None, idx2), edge_index, edge_type)
assert out2.size() == (2, 32)
assert torch.allclose(conv((idx1, idx2), edge_index, edge_type), out2)
assert conv((None, idx2), adj.t()).tolist() == out2.tolist()
assert conv((idx1, idx2), adj.t()).tolist() == out2.tolist()
t = '(Tuple[OptTensor, Tensor], Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index, edge_type).tolist() == out1.tolist()
if num_blocks is None:
assert torch.allclose(jit((None, idx2), edge_index, edge_type), out2)
assert torch.allclose(jit((idx1, idx2), edge_index, edge_type), out2)
t = '(Tuple[OptTensor, Tensor], SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
if num_blocks is None:
assert jit((None, idx2), adj.t()).tolist() == out2.tolist()
assert jit((idx1, idx2), adj.t()).tolist() == out2.tolist()
