def test_pdn_conv():
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
edge_attr = torch.randn(6, 8)
adj = SparseTensor(row=row, col=col, value=edge_attr, sparse_sizes=(4, 4))
conv = PDNConv(16, 32, edge_dim=8, hidden_channels=128)
assert str(conv) == "PDNConv(16, 32)"
out = conv(x, edge_index, edge_attr)
assert out.size() == (4, 32)
assert torch.allclose(conv(x, adj.t()), out, atol=1e-6)
t = '(Tensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x, edge_index, edge_attr), out)
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_pdn_conv_with_sparse_node_input_feature():
x = torch.sparse_coo_tensor(
indices=torch.tensor([[0, 0], [0, 1]]),
values=torch.tensor([1.0, 1.0]),
size=torch.Size([4, 16]),
)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
edge_attr = torch.randn(6, 8)
conv = PDNConv(16, 32, edge_dim=8, hidden_channels=128)
out = conv(x, edge_index, edge_attr)
assert out.size() == (4, 32)