def test_atomic_conv():
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
aconv = nn.AtomicConv(interaction_cutoffs=F.tensor([12.0, 12.0]),
rbf_kernel_means=F.tensor([0.0, 2.0]),
rbf_kernel_scaling=F.tensor([4.0, 4.0]),
features_to_use=F.tensor([6.0, 8.0]))
ctx = F.ctx()
if F.gpu_ctx():
aconv = aconv.to(ctx)
feat = F.randn((100, 1))
dist = F.randn((g.number_of_edges(), 1))
h = aconv(g, feat, dist)
# current we only do shape check
assert h.shape[-1] == 4
def test_atomic_conv(g, idtype):
g = g.astype(idtype).to(F.ctx())
aconv = nn.AtomicConv(interaction_cutoffs=F.tensor([12.0, 12.0]),
rbf_kernel_means=F.tensor([0.0, 2.0]),
rbf_kernel_scaling=F.tensor([4.0, 4.0]),
features_to_use=F.tensor([6.0, 8.0]))
ctx = F.ctx()
if F.gpu_ctx():
aconv = aconv.to(ctx)
feat = F.randn((g.number_of_nodes(), 1))
dist = F.randn((g.number_of_edges(), 1))
h = aconv(g, feat, dist)
# current we only do shape check
assert h.shape[-1] == 4
