def test_sddmm(g, shp, lhs_target, rhs_target, msg, index_dtype):
if dgl.backend.backend_name == 'mxnet' and g.number_of_edges() == 0:
pytest.skip() # mxnet do not support zero shape tensor
if dgl.backend.backend_name == 'tensorflow' and index_dtype == 'int32':
pytest.skip() # tensorflow dlpack has problem with int32 ndarray.
if index_dtype == 'int32':
g = g.int()
else:
g = g.long()
print(g)
print(g.idtype)
len_lhs = select(lhs_target, g.number_of_src_nodes(), g.number_of_edges(),
g.number_of_dst_nodes())
lhs_shp = (len_lhs, ) + shp[0]
len_rhs = select(rhs_target, g.number_of_src_nodes(), g.number_of_edges(),
g.number_of_dst_nodes())
rhs_shp = (len_rhs, ) + shp[1]
feat_lhs = F.tensor(np.random.rand(*lhs_shp) + 1)
feat_rhs = F.tensor(np.random.rand(*rhs_shp) + 1)
print('lhs shape: {}, rhs shape: {}'.format(F.shape(feat_lhs),
F.shape(feat_rhs)))
lhs_frame = select(lhs_target, g.srcdata, g.edata, g.dstdata)
rhs_frame = select(rhs_target, g.srcdata, g.edata, g.dstdata)
lhs_frame['x'] = F.attach_grad(F.clone(feat_lhs))
rhs_frame['y'] = F.attach_grad(F.clone(feat_rhs))
msg_func = lhs_target + '_' + msg + '_' + rhs_target
print('SDDMM(message func: {})'.format(msg_func))
lhs = F.attach_grad(F.clone(feat_lhs))
rhs = F.attach_grad(F.clone(feat_rhs))
with F.record_grad():
e = gsddmm(g,
msg,
lhs,
rhs,
lhs_target=lhs_target,
rhs_target=rhs_target)
F.backward(F.reduce_sum(e))
grad_lhs = F.grad(lhs)
grad_rhs = F.grad(rhs)
with F.record_grad():
g.apply_edges(udf_apply_edges[msg_func])
if g.number_of_edges() > 0:
e1 = g.edata['m']
assert F.allclose(e, e1)
print('forward passed')
F.backward(F.reduce_sum(e1))
if msg != 'copy_rhs':
assert F.allclose(F.grad(lhs_frame['x']), grad_lhs)
if msg != 'copy_lhs':
assert F.allclose(F.grad(rhs_frame['y']), grad_rhs)
print('backward passed')
lhs_frame.pop('x')
rhs_frame.pop('y')
if 'm' in g.edata: g.edata.pop('m')
def test_csrmm_backward(idtype, dtype, num_vtypes):
a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, 'A', 'B', 'AB')
b, B = _random_simple_graph(idtype, dtype, F.ctx(), 4, 3, 6, 'B', 'A' if num_vtypes == 1 else 'C', 'BA')
A_row, A_col = A.edges(order='eid')
B_row, B_col = B.edges(order='eid')
A_row = F.asnumpy(A_row)
A_col = F.asnumpy(A_col)
B_row = F.asnumpy(B_row)
B_col = F.asnumpy(B_col)
a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
b_dense = F.attach_grad(F.tensor(b.todense(), dtype=dtype))
A.edata['w'] = F.attach_grad(A.edata['w'])
B.edata['w'] = F.attach_grad(B.edata['w'])
with F.record_grad():
C = dgl.adj_product_graph(A, B, 'w')
assert len(C.ntypes) == num_vtypes
assert len(C.etypes) == 1
C_dense = np.zeros((3, 3))
C_row, C_col = C.edges(order='eid')
C_row = F.asnumpy(C_row)
C_col = F.asnumpy(C_col)
C_dense[C_row, C_col] = F.asnumpy(C.edata['w'])
c_dense = F.matmul(a_dense, b_dense)
assert np.allclose(C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4)
F.backward(F.reduce_sum(C.edata['w']) + F.reduce_sum(c_dense))
a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
b_dense_grad = F.asnumpy(F.grad(b_dense))[B_row, B_col]
A_spspmm_grad = F.asnumpy(F.grad(A.edata['w']))
B_spspmm_grad = F.asnumpy(F.grad(B.edata['w']))
assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
assert np.allclose(b_dense_grad, B_spspmm_grad, rtol=1e-4, atol=1e-4)
def test_edge_coarsening(idtype, g, weight, relabel):
num_nodes = g.num_nodes()
g = dgl.to_bidirected(g)
g = g.astype(idtype).to(F.ctx())
edge_weight = None
if weight:
edge_weight = F.abs(F.randn((g.num_edges(),))).to(F.ctx())
node_labels = neighbor_matching(g, edge_weight, relabel_idx=relabel)
unique_ids, counts = th.unique(node_labels, return_counts=True)
num_result_ids = unique_ids.size(0)
# shape correct
assert node_labels.shape == (g.num_nodes(),)
# all nodes marked
assert F.reduce_sum(node_labels < 0).item() == 0
# number of unique node ids correct.
assert num_result_ids >= num_nodes // 2 and num_result_ids <= num_nodes
# each unique id has <= 2 nodes
assert F.reduce_sum(counts > 2).item() == 0
# if two nodes have the same id, they must be neighbors
idxs = F.arange(0, num_nodes, idtype)
for l in unique_ids:
l = l.item()
idx = idxs[(node_labels == l)]
if idx.size(0) == 2:
u, v = idx[0].item(), idx[1].item()
assert g.has_edges_between(u, v)
def test_segment_reduce(reducer):
ctx = F.ctx()
value = F.tensor(np.random.rand(10, 5))
v1 = F.attach_grad(F.clone(value))
v2 = F.attach_grad(F.clone(value))
seglen = F.tensor([2, 3, 0, 4, 1, 0, 0])
u = F.copy_to(F.arange(0, F.shape(value)[0], F.int32), ctx)
v = F.repeat(F.copy_to(F.arange(0, len(seglen), F.int32), ctx),
seglen,
dim=0)
num_nodes = {'_U': len(u), '_V': len(seglen)}
g = dgl.convert.heterograph({('_U', '_E', '_V'): (u, v)},
num_nodes_dict=num_nodes)
with F.record_grad():
rst1 = gspmm(g, 'copy_lhs', reducer, v1, None)
if reducer in ['max', 'min']:
rst1 = F.replace_inf_with_zero(rst1)
F.backward(F.reduce_sum(rst1))
grad1 = F.grad(v1)
with F.record_grad():
rst2 = segment_reduce(seglen, v2, reducer=reducer)
F.backward(F.reduce_sum(rst2))
assert F.allclose(rst1, rst2)
print('forward passed')
grad2 = F.grad(v2)
assert F.allclose(grad1, grad2)
print('backward passed')
def test_csrmask_backward(idtype, dtype):
a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, 'A', 'B', 'AB')
b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, 'A', 'B', 'AB')
A_row, A_col = A.edges(order='eid')
B_row, B_col = B.edges(order='eid')
A_row = F.asnumpy(A_row)
A_col = F.asnumpy(A_col)
B_row = F.asnumpy(B_row)
B_col = F.asnumpy(B_col)
a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
A.edata['w'] = F.attach_grad(A.edata['w'])
with F.record_grad():
# Test for two element case
C1 = F.csrmask(A._graph, A.edata['w'], B._graph)
if dgl.backend.backend_name == 'tensorflow':
import tensorflow as tf
C2 = tf.gather_nd(a_dense, tf.stack([B_row, B_col], 1))
else:
C2 = a_dense[B_row, B_col]
assert F.allclose(C1, C2, rtol=1e-4, atol=1e-4)
F.backward(F.reduce_sum(C1) + F.reduce_sum(C2))
a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
A_spspmm_grad = F.asnumpy(F.grad(A.edata['w']))
assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
def test_edge_softmax(g, norm_by, shp, idtype):
g = g.astype(idtype).to(F.ctx())
edata = F.tensor(np.random.rand(g.number_of_edges(), *shp))
e1 = F.attach_grad(F.clone(edata))
with F.record_grad():
score1 = edge_softmax(g, e1, norm_by=norm_by)
F.backward(F.reduce_sum(score1))
grad_edata = F.grad(e1)
with F.record_grad():
e2 = F.attach_grad(F.clone(edata))
e2_2d = F.reshape(
e2,
(g.number_of_src_nodes(), g.number_of_dst_nodes(), *e2.shape[1:]))
if norm_by == 'src':
score2 = F.softmax(e2_2d, 1)
score2 = F.reshape(score2, (-1, *e2.shape[1:]))
if norm_by == 'dst':
score2 = F.softmax(e2_2d, 0)
score2 = F.reshape(score2, (-1, *e2.shape[1:]))
assert F.allclose(score1, score2)
print('forward passed')
F.backward(F.reduce_sum(score2))
assert F.allclose(F.grad(e2), grad_edata)
print('backward passed')
def _test(red, partial):
g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
# NOTE(zihao): add self-loop to avoid zero-degree nodes.
# https://github.com/dmlc/dgl/issues/761
g.add_edges(g.nodes(), g.nodes())
g = g.to(F.ctx())
hu, hv, he = generate_feature(g, 'none', 'none')
if partial:
nid = F.tensor(list(range(0, 100, 2)), g.idtype)
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_src(src='u', out='m'),
builtin[red](msg='m', out='r1'))
else:
g.update_all(fn.copy_src(src='u', out='m'),
builtin[red](msg='m', out='r1'))
r1 = g.ndata['r1']
F.backward(F.reduce_sum(r1))
n_grad1 = F.grad(g.ndata['u'])
# 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_src, udf_reduce[red])
else:
g.update_all(udf_copy_src, udf_reduce[red])
r2 = g.ndata['r2']
F.backward(F.reduce_sum(r2))
n_grad2 = F.grad(g.ndata['u'])
def _print_error(a, b):
print("ERROR: Test copy_src_{} partial: {}".
format(red, partial))
for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())):
if not np.allclose(x, y):
print('@{} {} v.s. {}'.format(i, x, y))
if not F.allclose(r1, r2):
_print_error(r1, r2)
assert F.allclose(r1, r2)
if F.backend_name != "jax":
if not F.allclose(n_grad1, n_grad2):
print('node grad')
_print_error(n_grad1, n_grad2)
assert(F.allclose(n_grad1, n_grad2))
def test_csrsum_backward(idtype, dtype, nelems):
a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, 'A', 'B', 'AB')
b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, 'A', 'B', 'AB')
A_row, A_col = A.edges(order='eid')
B_row, B_col = B.edges(order='eid')
A_row = F.asnumpy(A_row)
A_col = F.asnumpy(A_col)
B_row = F.asnumpy(B_row)
B_col = F.asnumpy(B_col)
a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
b_dense = F.attach_grad(F.tensor(b.todense(), dtype=dtype))
A.edata['w'] = F.attach_grad(A.edata['w'])
B.edata['w'] = F.attach_grad(B.edata['w'])
with F.record_grad():
if nelems == 2:
# Test for two element case
C = dgl.adj_sum_graph([A, B], 'w')
assert C.canonical_etypes == A.canonical_etypes
C_dense = np.zeros((3, 4))
C_row, C_col = C.edges(order='eid')
C_row = F.asnumpy(C_row)
C_col = F.asnumpy(C_col)
C_dense[C_row, C_col] = F.asnumpy(C.edata['w'])
c_dense = a_dense + b_dense
assert np.allclose(C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4)
F.backward(F.reduce_sum(C.edata['w']) + F.reduce_sum(c_dense))
a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
b_dense_grad = F.asnumpy(F.grad(b_dense))[B_row, B_col]
A_spspmm_grad = F.asnumpy(F.grad(A.edata['w']))
B_spspmm_grad = F.asnumpy(F.grad(B.edata['w']))
assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
assert np.allclose(b_dense_grad, B_spspmm_grad, rtol=1e-4, atol=1e-4)
elif nelems == 1:
# Test for single element case
C = dgl.adj_sum_graph([A], 'w')
assert C.canonical_etypes == A.canonical_etypes
C_dense = np.zeros((3, 4))
C_row, C_col = C.edges(order='eid')
C_row = F.asnumpy(C_row)
C_col = F.asnumpy(C_col)
C_dense[C_row, C_col] = F.asnumpy(C.edata['w'])
c_dense = a_dense
assert np.allclose(C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4)
F.backward(F.reduce_sum(C.edata['w']) + F.reduce_sum(c_dense))
a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
A_spspmm_grad = F.asnumpy(F.grad(A.edata['w']))
assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
def tensor_topo_traverse():
n = g.number_of_nodes()
mask = F.copy_to(F.ones((n, 1)), F.cpu())
degree = F.spmm(adjmat, mask)
while F.reduce_sum(mask) != 0.:
v = F.astype((degree == 0.), F.float32)
v = v * mask
mask = mask - v
frontier = F.copy_to(F.nonzero_1d(F.squeeze(v, 1)), F.cpu())
yield frontier
degree -= F.spmm(adjmat, v)
def test_spmm(idtype, g, shp, msg, reducer):
g = g.astype(idtype).to(F.ctx())
if dgl.backend.backend_name == 'tensorflow' and (reducer in ['min', 'max']):
pytest.skip() # tensorflow dlpack has problem writing into int32 arrays on GPU.
print(g)
print(g.idtype)
hu = F.tensor(np.random.rand(*((g.number_of_src_nodes(),) + shp[0])) + 1)
he = F.tensor(np.random.rand(*((g.number_of_edges(),) + shp[1])) + 1)
print('u shape: {}, e shape: {}'.format(F.shape(hu), F.shape(he)))
g.srcdata['x'] = F.attach_grad(F.clone(hu))
g.edata['w'] = F.attach_grad(F.clone(he))
print('SpMM(message func: {}, reduce func: {})'.format(msg, reducer))
u = F.attach_grad(F.clone(hu))
e = F.attach_grad(F.clone(he))
with F.record_grad():
v = gspmm(g, msg, reducer, u, e)
non_degree_indices = F.tensor(
np.nonzero(F.asnumpy(g.in_degrees()) != 0)[0])
v = F.gather_row(v, non_degree_indices)
if g.number_of_edges() > 0:
F.backward(F.reduce_sum(v))
if msg != 'copy_rhs':
grad_u = F.grad(u)
if msg != 'copy_lhs':
grad_e = F.grad(e)
with F.record_grad():
g.update_all(udf_msg[msg], udf_reduce[reducer])
if g.number_of_edges() > 0:
v1 = F.gather_row(g.dstdata['v'], non_degree_indices)
assert F.allclose(v, v1)
print('forward passed')
F.backward(F.reduce_sum(v1))
if msg != 'copy_rhs':
if reducer in ['min', 'max']: # there might be some numerical errors
rate = F.reduce_sum(F.abs(F.grad(g.srcdata['x']) - grad_u)) /\
F.reduce_sum(F.abs(grad_u))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.srcdata['x']), grad_u)
if msg != 'copy_lhs':
if reducer in ['min', 'max']:
rate = F.reduce_sum(F.abs(F.grad(g.edata['w']) - grad_e)) /\
F.reduce_sum(F.abs(grad_e))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.edata['w']), grad_e)
print('backward passed')
g.srcdata.pop('x')
g.edata.pop('w')
if 'v' in g.dstdata: g.dstdata.pop('v')
def test_spmm(idtype, g, shp, msg, reducer):
g = g.astype(idtype).to(F.ctx())
print(g)
print(g.idtype)
hu = F.tensor(np.random.rand(*((g.number_of_src_nodes(), ) + shp[0])) + 1)
he = F.tensor(np.random.rand(*((g.number_of_edges(), ) + shp[1])) + 1)
print('u shape: {}, e shape: {}'.format(F.shape(hu), F.shape(he)))
g.srcdata['x'] = F.attach_grad(F.clone(hu))
g.edata['w'] = F.attach_grad(F.clone(he))
print('SpMM(message func: {}, reduce func: {})'.format(msg, reducer))
u = F.attach_grad(F.clone(hu))
e = F.attach_grad(F.clone(he))
with F.record_grad():
v = gspmm(g, msg, reducer, u, e)
if reducer in ['max', 'min']:
v = F.replace_inf_with_zero(v)
if g.number_of_edges() > 0:
F.backward(F.reduce_sum(v))
if msg != 'copy_rhs':
grad_u = F.grad(u)
if msg != 'copy_lhs':
grad_e = F.grad(e)
with F.record_grad():
g.update_all(udf_msg[msg], udf_reduce[reducer])
if g.number_of_edges() > 0:
v1 = g.dstdata['v']
assert F.allclose(v, v1)
print('forward passed')
F.backward(F.reduce_sum(v1))
if msg != 'copy_rhs':
if reducer in ['min',
'max']: # there might be some numerical errors
rate = F.reduce_sum(F.abs(F.grad(g.srcdata['x']) - grad_u)) /\
F.reduce_sum(F.abs(grad_u))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.srcdata['x']), grad_u)
if msg != 'copy_lhs':
if reducer in ['min', 'max']:
rate = F.reduce_sum(F.abs(F.grad(g.edata['w']) - grad_e)) /\
F.reduce_sum(F.abs(grad_e))
assert F.as_scalar(rate) < 1e-2, rate
else:
assert F.allclose(F.grad(g.edata['w']), grad_e)
print('backward passed')
g.srcdata.pop('x')
g.edata.pop('w')
if 'v' in g.dstdata: g.dstdata.pop('v')
def _test(feat_scale):
in_feat = 16 * feat_scale
out_feat = 8 * feat_scale
print("in/out feat", in_feat, out_feat)
E_per_rel = F.copy_to(
F.tensor([
50, 100, 20, 284, 89, 10, 82, 9200, 10, 20, 30, 100, 128, 20,
284, 89, 10, 82, 92, 10, 20, 30, 100, 1280, 20, 284, 89, 1000,
82, 92, 10, 2000, 30, 100, 128, 20, 284, 89, 10, 82, 92, 10,
20, 30
]), F.cpu())
E_per_rel *= n_edge_scale
num_rel = len(E_per_rel)
print('num_rel', num_rel)
W_per_len = F.copy_to(
F.full((num_rel, ), in_feat, dtype=F.dtype(E_per_rel)), F.cpu())
H_arr = []
W_arr = []
Out_arr = []
Out_grad_arr = []
for eid in range(num_rel):
H_arr.append(F.randn((E_per_rel[eid], in_feat)))
W_arr.append(F.randn((in_feat, out_feat)))
Out_arr.append(F.zeros((E_per_rel[eid], out_feat)))
Out_grad_arr.append(F.ones((E_per_rel[eid], out_feat)))
H = F.cat([h for h in H_arr], 0)
W = F.cat([w for w in W_arr], 0)
W_3D = W.reshape(num_rel, in_feat, out_feat)
Out = F.cat([out for out in Out_arr], 0)
Out_grad = F.cat([o for o in Out_grad_arr], 0)
print('H.shape', H.shape)
print('W.shape', W.shape)
print('W_3D.shape', W_3D.shape)
print('Out.shape', Out.shape)
etype_arr = []
for eid in range(num_rel):
etype_arr.append(
F.full((E_per_rel[eid], ), eid, dtype=F.dtype(E_per_rel)))
etypes = F.cat([etype for etype in etype_arr], 0)
#################################################################
# low-mem version using PyTorch operator
#################################################################
# forward pass
out = []
for i in range(len(E_per_rel)):
Hi = H_arr[i]
Wi = W_arr[i]
out.append(F.matmul(Hi, Wi))
out_low_mem = F.cat(out, 0)
# backward pass
H_grad = []
W_grad = []
for i in range(len(E_per_rel)):
Hi = H_arr[i]
Wi = W_arr[i]
Out_gradi = Out_grad_arr[i]
H_grad.append(F.matmul(Out_gradi, Wi.transpose(0, 1)))
W_grad.append(F.matmul(Hi.transpose(0, 1), Out_gradi))
Hgrad_low_mem = F.cat(H_grad, 0)
Wgrad_low_mem = F.cat(W_grad, 0)
Wgrad_low_mem = Wgrad_low_mem.reshape(num_rel, in_feat, out_feat)
#################################################################
# gather_mm where H sorted according to etype
#################################################################
seglen_A = E_per_rel
F.attach_grad(H)
F.attach_grad(W_3D)
with F.record_grad():
out_gmm_sorted = dgl.ops.segment_mm(H, W_3D, seglen_A)
F.backward(F.reduce_sum(out_gmm_sorted))
Hgrad_gmm_sorted = H.grad
Wgrad_gmm_sorted = W_3D.grad
#################################################################
# gather_mm where H is not sorted (backward not supported yet)
#################################################################
F.attach_grad(H)
F.attach_grad(W_3D)
with F.record_grad():
out_gmm_unsorted = dgl.ops.gather_mm(H, W_3D, idx_rhs=etypes)
F.backward(F.reduce_sum(out_gmm_unsorted))
Hgrad_gmm_unsorted = H.grad
Wgrad_gmm_unsorted = W_3D.grad
# correctness check
assert F.allclose(out_low_mem, out_gmm_sorted, atol=1e-3, rtol=1e-3)
assert F.allclose(Hgrad_low_mem,
Hgrad_gmm_sorted,
atol=1e-3,
rtol=1e-3)
assert F.allclose(Wgrad_low_mem,
Wgrad_gmm_sorted,
atol=1e-3,
rtol=1e-3)
assert F.allclose(out_low_mem, out_gmm_unsorted, atol=1e-3, rtol=1e-3)
assert F.allclose(Hgrad_low_mem,
Hgrad_gmm_unsorted,
atol=1e-3,
rtol=1e-3)
assert F.allclose(Wgrad_low_mem,
Wgrad_gmm_unsorted,
atol=1e-3,
rtol=1e-3)
def _test(g, lhs, rhs, binary_op, reducer, partial, nid, broadcast='none'):
# initialize node/edge features with uniform(-1, 1)
hu, hv, he = generate_feature(g, broadcast, binary_op)
if binary_op == 'div':
# op = div
# lhs range: [-1, 1]
# rhs range: [1, 2]
# result range: [-1, 1]
if rhs == 'u':
hu = (hu + 3) / 2
elif rhs == 'v':
hv = (hv + 3) / 2
elif rhs == 'e':
he = (he + 3) / 2
if binary_op == 'add' or binary_op == 'sub':
# op = add, sub
# lhs range: [-1/2, 1/2]
# rhs range: [-1/2, 1/2]
# result range: [-1, 1]
hu = hu / 2
hv = hv / 2
he = he / 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))
builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs)
builtin_msg = getattr(fn, builtin_msg_name)
builtin_red = getattr(fn, reducer)
def target_feature_switch(g, target):
if target == "u":
return g.ndata["u"]
elif target == "v":
return g.ndata["v"]
else:
return g.edata["e"]
with F.record_grad():
if partial:
g.pull(nid, builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1'))
else:
g.update_all(builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1'))
r1 = g.ndata.pop('r1')
F.backward(F.reduce_sum(r1))
lhs_grad_1 = F.grad(target_feature_switch(g, lhs))
rhs_grad_1 = F.grad(target_feature_switch(g, rhs))
# 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))
def target_switch(edges, target):
if target == "u":
return edges.src
elif target == "v":
return edges.dst
elif target == "e":
return edges.data
else:
assert(0), "Unknown target {}".format(target)
def mfunc(edges):
op = getattr(F, binary_op)
lhs_data = target_switch(edges, lhs)[lhs]
rhs_data = target_switch(edges, rhs)[rhs]
# NOTE(zihao): we need to do batched broadcast
# e.g. (68, 3, 1) op (68, 5, 3, 4)
while F.ndim(lhs_data) < F.ndim(rhs_data):
lhs_data = F.unsqueeze(lhs_data, 1)
while F.ndim(rhs_data) < F.ndim(lhs_data):
rhs_data = F.unsqueeze(rhs_data, 1)
return {"m": op(lhs_data, rhs_data)}
def rfunc(nodes):
op = getattr(F, reducer)
return {"r2": op(nodes.mailbox['m'], 1)}
with F.record_grad():
if partial:
g.pull(nid, mfunc, rfunc)
else:
g.update_all(mfunc, rfunc)
r2 = g.ndata.pop('r2')
F.backward(F.reduce_sum(r2), F.tensor([1.]))
lhs_grad_2 = F.grad(target_feature_switch(g, lhs))
rhs_grad_2 = F.grad(target_feature_switch(g, rhs))
rtol = 1e-4
atol = 1e-4
def _print_error(a, b):
print("ERROR: Test {}_{}_{}_{} broadcast: {} partial: {}".
format(lhs, binary_op, rhs, reducer, broadcast, partial))
return
if lhs == 'u':
lhs_data = hu
elif lhs == 'v':
lhs_data = hv
elif lhs == 'e':
lhs_data = he
if rhs == 'u':
rhs_data = hu
elif rhs == 'v':
rhs_data = hv
elif rhs == 'e':
rhs_data = he
print("lhs", F.asnumpy(lhs_data).tolist())
print("rhs", F.asnumpy(rhs_data).tolist())
for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())):
if not np.allclose(x, y, rtol, atol):
print('@{} {} v.s. {}'.format(i, x, y))
if not F.allclose(r1, r2, rtol, atol):
_print_error(r1, r2)
assert F.allclose(r1, r2, rtol, atol)
if not F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol):
print("left grad")
_print_error(lhs_grad_1, lhs_grad_2)
assert(F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol))
if not F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol):
print("right grad")
_print_error(rhs_grad_1, rhs_grad_2)
assert(F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol))
def test_edge_softmax(g, norm_by, idtype):
print("params", norm_by, idtype)
g = create_test_heterograph(idtype)
x1 = F.randn((g.num_edges('plays'), feat_size))
x2 = F.randn((g.num_edges('follows'), feat_size))
x3 = F.randn((g.num_edges('develops'), feat_size))
x4 = F.randn((g.num_edges('wishes'), feat_size))
F.attach_grad(F.clone(x1))
F.attach_grad(F.clone(x2))
F.attach_grad(F.clone(x3))
F.attach_grad(F.clone(x4))
g['plays'].edata['eid'] = x1
g['follows'].edata['eid'] = x2
g['develops'].edata['eid'] = x3
g['wishes'].edata['eid'] = x4
#################################################################
# edge_softmax() on homogeneous graph
#################################################################
with F.record_grad():
hm_g = dgl.to_homogeneous(g)
hm_x = F.cat((x3, x2, x1, x4), 0)
hm_e = F.attach_grad(F.clone(hm_x))
score_hm = edge_softmax(hm_g, hm_e, norm_by=norm_by)
hm_g.edata['score'] = score_hm
ht_g = dgl.to_heterogeneous(hm_g, g.ntypes, g.etypes)
r1 = ht_g.edata['score'][('user', 'plays', 'game')]
r2 = ht_g.edata['score'][('user', 'follows', 'user')]
r3 = ht_g.edata['score'][('developer', 'develops', 'game')]
r4 = ht_g.edata['score'][('user', 'wishes', 'game')]
F.backward(F.reduce_sum(r1) + F.reduce_sum(r2))
grad_edata_hm = F.grad(hm_e)
#################################################################
# edge_softmax() on heterogeneous graph
#################################################################
e1 = F.attach_grad(F.clone(x1))
e2 = F.attach_grad(F.clone(x2))
e3 = F.attach_grad(F.clone(x3))
e4 = F.attach_grad(F.clone(x4))
e = {
('user', 'follows', 'user'): e2,
('user', 'plays', 'game'): e1,
('user', 'wishes', 'game'): e4,
('developer', 'develops', 'game'): e3
}
with F.record_grad():
score = edge_softmax(g, e, norm_by=norm_by)
r5 = score[('user', 'plays', 'game')]
r6 = score[('user', 'follows', 'user')]
r7 = score[('developer', 'develops', 'game')]
r8 = score[('user', 'wishes', 'game')]
F.backward(F.reduce_sum(r5) + F.reduce_sum(r6))
grad_edata_ht = F.cat((F.grad(e3), F.grad(e2), F.grad(e1), F.grad(e4)),
0)
# correctness check
assert F.allclose(r1, r5)
assert F.allclose(r2, r6)
assert F.allclose(r3, r7)
assert F.allclose(r4, r8)
assert F.allclose(grad_edata_hm, grad_edata_ht)
