def verify_graph_feats(g, part, node_feats):
for ntype in g.ntypes:
ntype_id = g.get_ntype_id(ntype)
for name in g.nodes[ntype].data:
if name in [dgl.NID, 'inner_node']:
continue
inner_node_mask = _get_inner_node_mask(part, ntype_id)
inner_nids = F.boolean_mask(part.ndata[dgl.NID], inner_node_mask)
min_nids = F.min(inner_nids, 0)
orig_id = F.boolean_mask(part.ndata['orig_id'], inner_node_mask)
true_feats = F.gather_row(g.nodes[ntype].data[name], orig_id)
ndata = F.gather_row(node_feats[ntype + '/' + name],
inner_nids - min_nids)
assert np.all(F.asnumpy(ndata == true_feats))
return foo
udf_apply_edges = {
lhs_target + '_' + msg + '_' + rhs_target:
edge_func(lhs_target, rhs_target, msg)
for lhs_target in ['u', 'v', 'e'] for rhs_target in ['u', 'v', 'e']
for msg in ['add', 'sub', 'mul', 'div', 'dot', 'copy_lhs', 'copy_rhs']
}
udf_reduce = {
'sum': lambda nodes: {
'v': F.sum(nodes.mailbox['m'], 1)
},
'min': lambda nodes: {
'v': F.min(nodes.mailbox['m'], 1)
},
'max': lambda nodes: {
'v': F.max(nodes.mailbox['m'], 1)
}
}
graphs = [
# dgl.rand_graph(30, 0),
dgl.rand_graph(100, 30),
dgl.rand_graph(100, 3000),
dgl.rand_bipartite(80, 160, 3000)
]
spmm_shapes = [((1, 2, 1, 3, 1), (4, 1, 3, 1, 1)),
((5, 3, 1, 7), (1, 3, 7, 1)), ((1, 3, 1), (4, 1, 3)),
select(lhs_target, edges.src, edges.data, edges.dst)['x'],
select(rhs_target, edges.src, edges.data, edges.dst)['y']
)
}
return foo
udf_apply_edges = {
lhs_target + '_' + msg + '_' + rhs_target: edge_func(lhs_target, rhs_target, msg)
for lhs_target in ['u', 'v', 'e']
for rhs_target in ['u', 'v', 'e']
for msg in ['add', 'sub', 'mul', 'div', 'dot', 'copy_lhs', 'copy_rhs']
}
udf_reduce = {
'sum': lambda nodes: {'v': F.sum(nodes.mailbox['m'], 1)},
'min': lambda nodes: {'v': F.min(nodes.mailbox['m'], 1)},
'max': lambda nodes: {'v': F.max(nodes.mailbox['m'], 1)}
}
graphs = [
# dgl.rand_graph(30, 0),
dgl.rand_graph(30, 100),
dgl.rand_bipartite('_U', '_E', '_V', 30, 40, 300)
]
spmm_shapes = [
((1, 2, 1, 3, 1), (4, 1, 3, 1, 1)),
((3, 3), (1, 3)),
((1,), (3,)),
((3,), (1,)),
((1,), (1,)),