def test_appnp():
x = torch.randn(4, 16)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = APPNP(K=10, alpha=0.1)
assert conv.__repr__() == 'APPNP(K=10, alpha=0.1)'
out = conv(x, edge_index)
assert out.size() == (4, 16)
assert torch.allclose(conv(x, adj.t()), out, atol=1e-6)
t = '(Tensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x, edge_index).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x, adj.t()), out, atol=1e-6)
def test_appnp():
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.rand(edge_index.size(1))
num_nodes = edge_index.max().item() + 1
x = torch.randn((num_nodes, in_channels))
lin = torch.nn.Linear(in_channels, out_channels)
conv = APPNP(K=10, alpha=0.1)
assert conv.__repr__() == 'APPNP(K=10, alpha=0.1)'
out1 = conv(lin(x), edge_index)
assert out1.size() == (num_nodes, out_channels)
out2 = conv(lin(x), edge_index, edge_weight)
assert out2.size() == (num_nodes, out_channels)
jit = torch.jit.script(conv.jittable())
assert jit(lin(x), edge_index).tolist() == out1.tolist()
assert jit(lin(x), edge_index, edge_weight).tolist() == out2.tolist()
