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_feast_conv():
in_channels, out_channels = (16, 32)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
num_nodes = edge_index.max().item() + 1
x = torch.randn((num_nodes, in_channels))
conv = FeaStConv(in_channels, out_channels, heads=2)
assert conv.__repr__() == 'FeaStConv(16, 32, heads=2)'
out = conv(x, edge_index)
assert out.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() == out.tolist()
