Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def forward(self, g, lg, x, y, deg_g, deg_lg, pm_pd):
pmpd_x = F.embedding(pm_pd, x)
sum_x = sum(theta(z) for theta, z in zip(self.theta_list, self.aggregate(g, x)))
g.set_e_repr({'y' : y})
g.update_all(fn.copy_edge(edge='y', out='m'), fn.sum('m', 'pmpd_y'))
pmpd_y = g.pop_n_repr('pmpd_y')
x = self.theta_x(x) + self.theta_deg(deg_g * x) + sum_x + self.theta_y(pmpd_y)
n = self.out_feats // 2
x = th.cat([x[:, :n], F.relu(x[:, n:])], 1)
x = self.bn_x(x)
sum_y = sum(gamma(z) for gamma, z in zip(self.gamma_list, self.aggregate(lg, y)))
y = self.gamma_y(y) + self.gamma_deg(deg_lg * y) + sum_y + self.gamma_x(pmpd_x)
y = th.cat([y[:, :n], F.relu(y[:, n:])], 1)
y = self.bn_y(y)
return x, y
col.append(dst)
etypes.append(th.from_numpy(etype))
# update shift
v_shift += g.number_of_nodes
e_shift += g.number_of_edges
n = v_shift
leaf_ids = th.cat(leaf_ids)
readout_ids = th.cat(readout_ids)
pos_arr = th.cat(pos_arr)
etypes = th.cat(etypes)
row, col = map(np.concatenate, (row, col))
coo = coo_matrix((np.zeros_like(row), (row, col)), shape=(n, n))
g = dgl.DGLGraph(coo, readonly=True)
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
data = th.cat(data)
labels = th.cat(labels)
g.edata['etype'] = etypes
if self.graph_type == 'openai':
g.ndata['pos_0'] = pos_arr // self.attrs['stride']
g.ndata['pos_1'] = pos_arr % self.attrs['stride']
else:
g.ndata['pos'] = pos_arr
g.nodes[leaf_ids].data['x'] = data
return Batch(g=g, readout_ids=readout_ids, leaf_ids=leaf_ids, y=labels)
row.append(src)
col.append(dst)
etypes.append(th.from_numpy(etype))
# update shift
v_shift += g.number_of_nodes
e_shift += g.number_of_edges
n = v_shift
leaf_ids = th.cat(leaf_ids)
readout_ids = th.cat(readout_ids)
pos_arr = th.cat(pos_arr)
etypes = th.cat(etypes)
row, col = map(np.concatenate, (row, col))
coo = coo_matrix((np.zeros_like(row), (row, col)), shape=(n, n))
g = dgl.DGLGraph(coo, readonly=True)
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
data = th.cat(data)
labels = th.cat(labels)
g.edata['etype'] = etypes
if self.graph_type == 'openai':
g.ndata['pos_0'] = pos_arr // self.attrs['stride']
g.ndata['pos_1'] = pos_arr % self.attrs['stride']
else:
g.ndata['pos'] = pos_arr
g.nodes[leaf_ids].data['x'] = data
return Batch(g=g, readout_ids=readout_ids, leaf_ids=leaf_ids, y=labels)
def _test(apply_func):
g = generate_graph()
f = g.ndata['f']
# an out place run to get result
g.pull(nodes,
fn.copy_src(src='f', out='m'), fn.sum(msg='m', out='f'), apply_func)
result = g.ndata['f']
# inplace deg bucket
v1 = F.clone(f)
g.ndata['f'] = v1
g.pull(nodes, message_func, reduce_func, apply_func, inplace=True)
r1 = g.ndata['f']
# check result
assert F.allclose(r1, result)
# check inplace
assert F.allclose(v1, r1)
# inplace v2v spmv
v1 = F.clone(f)
g.ndata['f'] = v1
g.pull(nodes, fn.copy_src(src='f', out='m'),
def _test(red, partial):
g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
# NOTE(zihao): add self-loop to avoid zero-degree nodes.
g.add_edges(g.nodes(), g.nodes())
hu, hv, he = generate_feature(g, 'none', 'none')
if partial:
nid = F.tensor(list(range(0, 100, 2)))
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
with F.record_grad():
if partial:
g.pull(nid, fn.copy_edge(edge='e', out='m'),
builtin[red](msg='m', out='r1'))
else:
g.update_all(fn.copy_edge(edge='e', out='m'),
builtin[red](msg='m', out='r1'))
r1 = g.ndata['r1']
F.backward(F.reduce_sum(r1))
e_grad1 = F.grad(g.edata['e'])
# reset grad
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
with F.record_grad():
if partial:
g.pull(nid, udf_copy_edge, udf_reduce[red])
def test_edge_batch():
d = 10
g = dgl.DGLGraph(nx.path_graph(20))
nfeat = F.randn((g.number_of_nodes(), d))
efeat = F.randn((g.number_of_edges(), d))
g.ndata['x'] = nfeat
g.edata['x'] = efeat
# test all
eid = utils.toindex(slice(0, g.number_of_edges()))
u, v, _ = g._graph.edges('eid')
src_data = g.get_n_repr(u)
edge_data = g.get_e_repr(eid)
dst_data = g.get_n_repr(v)
ebatch = EdgeBatch((u, v, eid), src_data, edge_data, dst_data)
assert F.shape(ebatch.src['x'])[0] == g.number_of_edges() and\
F.shape(ebatch.src['x'])[1] == d
assert F.shape(ebatch.dst['x'])[0] == g.number_of_edges() and\
def test_local_var():
g = DGLGraph(nx.path_graph(5))
g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
g.edata['w'] = F.zeros((g.number_of_edges(), 4))
# test override
def foo(g):
g = g.local_var()
g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
g.edata['w'] = F.ones((g.number_of_edges(), 4))
foo(g)
assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
# test out-place update
def foo(g):
g = g.local_var()
g.nodes[[2, 3]].data['h'] = F.ones((2, 3))
g.edges[[2, 3]].data['w'] = F.ones((2, 4))
foo(g)
def test_spmv_3d_feat():
def src_mul_edge_udf(edges):
return {'sum': edges.src['h'] * F.unsqueeze(F.unsqueeze(edges.data['h'], 1), 1)}
def sum_udf(nodes):
return {'h': F.sum(nodes.mailbox['sum'], 1)}
n = 100
p = 0.1
a = sp.random(n, n, p, data_rvs=lambda n: np.ones(n))
g = dgl.DGLGraph(a)
m = g.number_of_edges()
# test#1: v2v with adj data
h = F.randn((n, 5, 5))
e = F.randn((m,))
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=fn.src_mul_edge('h', 'h', 'sum'), reduce_func=fn.sum('sum', 'h')) # 1
ans = g.ndata['h']
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=fn.sum('sum', 'h')) # 2
assert F.allclose(g.ndata['h'], ans)
def test_batch_no_edge():
g1 = dgl.DGLGraph()
g1.add_nodes(6)
g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
g2 = dgl.DGLGraph()
g2.add_nodes(6)
g2.add_edges([0, 1, 2, 5, 4, 5], [1 ,2 ,3, 4, 3, 0])
g3 = dgl.DGLGraph()
g3.add_nodes(1) # no edges
g = dgl.batch([g1, g3, g2]) # should not throw an error
def generate_graph():
g = DGLGraph()
for i in range(10):
g.add_node(i, __REPR__=i+1) # 10 nodes.
# create a graph where 0 is the source and 9 is the sink
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
return g