def test_softmax_edges():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(10))
feat0 = F.randn((g0.number_of_edges(), 10))
g0.edata['x'] = feat0
ground_truth = F.softmax(feat0, dim=0)
assert F.allclose(dgl.softmax_edges(g0, 'x'), ground_truth)
g0.edata.pop('x')
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(5))
g2 = dgl.DGLGraph(nx.path_graph(3))
g3 = dgl.DGLGraph()
g4 = dgl.DGLGraph(nx.path_graph(10))
bg = dgl.batch([g0, g1, g2, g3, g4])
feat1 = F.randn((g1.number_of_edges(), 10))
feat2 = F.randn((g2.number_of_edges(), 10))
feat4 = F.randn((g4.number_of_edges(), 10))
bg.edata['x'] = F.cat([feat0, feat1, feat2, feat4], 0)
ground_truth = F.cat([
F.softmax(feat0, 0),
F.softmax(feat1, 0),
F.softmax(feat2, 0),
F.softmax(feat4, 0)
], 0)
assert F.allclose(dgl.softmax_edges(bg, 'x'), ground_truth)
def test_softmax(g, idtype):
g = g.astype(idtype).to(F.ctx())
g.ndata['h'] = F.randn((g.number_of_nodes(), 3))
g.edata['h'] = F.randn((g.number_of_edges(), 2))
# Test.1: node readout
x = dgl.softmax_nodes(g, 'h')
subg = dgl.unbatch(g)
subx = []
for sg in subg:
subx.append(F.softmax(sg.ndata['h'], dim=0))
assert F.allclose(x, F.cat(subx, dim=0))
# Test.2: edge readout
x = dgl.softmax_edges(g, 'h')
subg = dgl.unbatch(g)
subx = []
for sg in subg:
subx.append(F.softmax(sg.edata['h'], dim=0))
assert F.allclose(x, F.cat(subx, dim=0))