def test_dna_conv():
channels = 32
num_layers = 3
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, num_layers, channels))
conv = DNAConv(channels, heads=4, groups=8, dropout=0.0)
assert conv.__repr__() == 'DNAConv(32, heads=4, groups=8)'
out = conv(x, edge_index)
assert out.size() == (num_nodes, channels)
if is_full_test():
jit = torch.jit.script(conv.jittable())
assert jit(x, edge_index).tolist() == out.tolist()
conv = DNAConv(channels, heads=1, groups=1, dropout=0.0)
assert conv.__repr__() == 'DNAConv(32, heads=1, groups=1)'
out = conv(x, edge_index)
assert out.size() == (num_nodes, channels)
if is_full_test():
jit = torch.jit.script(conv.jittable())
assert jit(x, edge_index).tolist() == out.tolist()
conv = DNAConv(channels, heads=1, groups=1, dropout=0.0, cached=True)
out = conv(x, edge_index)
out = conv(x, edge_index)
assert out.size() == (num_nodes, channels)
if is_full_test():
jit = torch.jit.script(conv.jittable())
assert jit(x, edge_index).tolist() == out.tolist()
def test_dna_conv():
channels = 32
num_layers = 3
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, num_layers, channels))
conv = DNAConv(channels, heads=4, groups=8, dropout=0.5)
assert conv.__repr__() == 'DNAConv(32, heads=4, groups=8)'
assert conv(x, edge_index).size() == (num_nodes, channels)
conv = DNAConv(channels, heads=1, groups=1, dropout=0.5)
assert conv.__repr__() == 'DNAConv(32, heads=1, groups=1)'
assert conv(x, edge_index).size() == (num_nodes, channels)
def test_dna_conv():
x = torch.randn((4, 3, 32))
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
row, col = edge_index
value = torch.rand(row.size(0))
adj2 = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
adj1 = adj2.set_value(None)
conv = DNAConv(32, heads=4, groups=8, dropout=0.0)
assert conv.__repr__() == 'DNAConv(32, heads=4, groups=8)'
out1 = conv(x, edge_index)
assert out1.size() == (4, 32)
assert torch.allclose(conv(x, adj1.t()), out1, atol=1e-6)
out2 = conv(x, edge_index, value)
assert out2.size() == (4, 32)
assert torch.allclose(conv(x, adj2.t()), out2, atol=1e-6)
if is_full_test():
t = '(Tensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x, edge_index).tolist() == out1.tolist()
assert jit(x, edge_index, value).tolist() == out2.tolist()
t = '(Tensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x, adj1.t()), out1, atol=1e-6)
assert torch.allclose(jit(x, adj2.t()), out2, atol=1e-6)
conv.cached = True
conv(x, edge_index)
assert conv(x, edge_index).tolist() == out1.tolist()
conv(x, adj1.t())
assert torch.allclose(conv(x, adj1.t()), out1, atol=1e-6)