Skip to content

GitLab

Unverified Commit 44089c8b authored by Minjie Wang's avatar Minjie Wang Committed by GitHub
Browse files

[Refactor][Graph] Merge DGLGraph and DGLHeteroGraph (#1862) 



* Merge

* [Graph][CUDA] Graph on GPU and many refactoring (#1791)

* change edge_ids behavior and C++ impl

* fix unittests; remove utils.Index in edge_id

* pass mx and th tests

* pass tf test

* add aten::Scatter_

* Add nonzero; impl CSRGetDataAndIndices/CSRSliceMatrix

* CSRGetData and CSRGetDataAndIndices passed tests

* CSRSliceMatrix basic tests

* fix bug in empty slice

* CUDA CSRHasDuplicate

* has_node; has_edge_between

* predecessors, successors

* deprecate send/recv; fix send_and_recv

* deprecate send/recv; fix send_and_recv

* in_edges; out_edges; all_edges; apply_edges

* in deg/out deg

* subgraph/edge_subgraph

* adj

* in_subgraph/out_subgraph

* sample neighbors

* set/get_n/e_repr

* wip: working on refactoring all idtypes

* pass ndata/edata tests on gpu

* fix

* stash

* workaround nonzero issue

* stash

* nx conversion

* test_hetero_basics except update routines

* test_update_routines

* test_hetero_basics for pytorch

* more fixes

* WIP: flatten graph

* wip: flatten

* test_flatten

* test_to_device

* fix bug in to_homo

* fix bug in CSRSliceMatrix

* pass subgraph test

* fix send_and_recv

* fix filter

* test_heterograph

* passed all pytorch tests

* fix mx unittest

* fix pytorch test_nn

* fix all unittests for PyTorch

* passed all mxnet tests

* lint

* fix tf nn test

* pass all tf tests

* lint

* lint

* change deprecation

* try fix compile

* lint

* update METIDS

* fix utest

* fix

* fix utests

* try debug

* revert

* small fix

* fix utests

* upd

* upd

* upd

* fix

* upd

* upd

* upd

* upd

* upd

* trigger

* +1s

* [kernel] Use heterograph index instead of unitgraph index (#1813)

* upd

* upd

* upd

* fix

* upd

* upd

* upd

* upd

* upd

* trigger

* +1s

* [Graph] Mutation for Heterograph (#1818)

* mutation add_nodes and add_edges

* Add support for remove_edges, remove_nodes, add_selfloop, remove_selfloop

* Fix
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-51-214.ec2.internal>

* upd

* upd

* upd

* fix

* [Transfom] Mutable transform (#1833)

* add nodesy

* All three

* Fix

* lint

* Add some test case

* Fix

* Fix

* Fix

* Fix

* Fix

* Fix

* fix

* triger

* Fix

* fix
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-51-214.ec2.internal>

* [Graph] Migrate Batch & Readout module to heterograph (#1836)

* dgl.batch

* unbatch

* fix to device

* reduce readout; segment reduce

* change batch_num_nodes|edges to function

* reduce readout/ softmax

* broadcast

* topk

* fix

* fix tf and mx

* fix some ci

* fix batch but unbatch differently

* new checkk

* upd

* upd

* upd

* idtype behavior; code reorg

* idtype behavior; code reorg

* wip: test_basics

* pass test_basics

* WIP: from nx/ to nx

* missing files

* upd

* pass test_basics:test_nx_conversion

* Fix test

* Fix inplace update

* WIP: fixing tests

* upd

* pass test_transform cpu

* pass gpu test_transform

* pass test_batched_graph

* GPU graph auto cast to int32

* missing file

* stash

* WIP: rgcn-hetero

* Fix two datasety

* upd

* weird

* Fix capsuley

* fuck you

* fuck matthias

* Fix dgmg

* fix bug in block degrees; pass rgcn-hetero

* rgcn

* gat and diffpool fix
also fix ppi and tu dataset

* Tree LSTM

* pointcloud

* rrn; wip: sgc

* resolve conflicts

* upd

* sgc and reddit dataset

* upd

* Fix deepwalk, gindt and gcn

* fix datasets and sign

* optimization

* optimization

* upd

* upd

* Fix GIN

* fix bug in add_nodes add_edges; tagcn

* adaptive sampling and gcmc

* upd

* upd

* fix geometric

* fix

* metapath2vec

* fix agnn

* fix pickling problem of block

* fix utests

* miss file

* linegraph

* upd

* upd

* upd

* graphsage

* stgcn_wave

* fix hgt

* on unittests

* Fix transformer

* Fix HAN

* passed pytorch unittests

* lint

* fix

* Fix cluster gcn

* cluster-gcn is ready

* on fixing block related codes

* 2nd order derivative

* Revert "2nd order derivative"

This reverts commit 523bf6c249bee61b51b1ad1babf42aad4167f206.

* passed torch utests again

* fix all mxnet unittests

* delete some useless tests

* pass all tf cpu tests

* disable

* disable distributed unittest

* fix

* fix

* lint

* fix

* fix

* fix script

* fix tutorial

* fix apply edges bug

* fix 2 basics

* fix tutorial
Co-authored-by: default avataryzh119 <expye@outlook.com>
Co-authored-by: default avatarxiang song(charlie.song) <classicxsong@gmail.com>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-51-214.ec2.internal>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-7-42.us-west-2.compute.internal>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-1-5.us-west-2.compute.internal>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-68-185.ec2.internal>
parent 015acfd2
Hide whitespace changes
Inline Side-by-side
Showing with 1048 additions and 1232 deletions
tests/compute/test_kernel.py
View file @ 44089c8b
......@@ -71,7 +71,7 @@ def generate_feature(g, broadcast='none', binary_op='none'):
u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
return u, v, e
return F.astype(u, F.float32), F.astype(v, F.float32), F.astype(e, F.float32)
def test_copy_src_reduce():
......@@ -80,9 +80,10 @@ def test_copy_src_reduce():
# 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)))
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))
......@@ -141,9 +142,10 @@ def test_copy_edge_reduce():
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())
g = g.to(F.ctx())
hu, hv, he = generate_feature(g, 'none', 'none')
if partial:
nid = F.tensor(list(range(0, 100, 2)))
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))
......@@ -348,7 +350,8 @@ def test_all_binary_builtins():
g.add_edge(18, 1)
g.add_edge(19, 0)
g.add_edge(19, 1)
nid = F.tensor([0, 1, 4, 5, 7, 12, 14, 15, 18, 19])
g = g.to(F.ctx())
nid = F.tensor([0, 1, 4, 5, 7, 12, 14, 15, 18, 19], g.idtype)
target = ["u", "v", "e"]
for lhs, rhs in product(target, target):
......
tests/compute/test_multi_send_recv.py deleted 100644 → 0
View file @ 015acfd2
import numpy as np
import dgl
from dgl.graph import DGLGraph
from collections import defaultdict as ddict
import scipy.sparse as sp
import backend as F
D = 5
def message_func(edges):
assert len(edges.src['h'].shape) == 2
assert edges.src['h'].shape[1] == D
return {'m' : edges.src['h']}
def reduce_func(nodes):
msgs = nodes.mailbox['m']
assert len(msgs.shape) == 3
assert msgs.shape[2] == D
return {'accum' : F.sum(msgs, 1)}
def apply_node_func(nodes):
return {'h' : nodes.data['h'] + nodes.data['accum']}
def generate_graph(grad=False):
g = DGLGraph()
g.add_nodes(10) # 10 nodes.
# create a graph where 0 is the source and 9 is the sink
# 16 edges
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
ncol = F.randn((10, D))
ecol = F.randn((16, D))
if grad:
ncol = F.attach_grad(ncol)
ecol = F.attach_grad(ecol)
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
g.ndata['h'] = ncol
g.edata['w'] = ecol
return g
def test_multi_send():
g = generate_graph()
def _fmsg(edges):
assert edges.src['h'].shape == (5, D)
return {'m' : edges.src['h']}
g.register_message_func(_fmsg)
# many-many send
u = F.tensor([0, 0, 0, 0, 0])
v = F.tensor([1, 2, 3, 4, 5])
g.send((u, v))
# duplicate send
u = F.tensor([0])
v = F.tensor([1, 2, 3, 4, 5])
g.send((u, v))
# send more
u = F.tensor([1, 2, 3, 4, 5])
v = F.tensor([9])
g.send((u, v))
# check if message indicator is as expected
expected = F.copy_to(F.zeros((g.number_of_edges(),), dtype=F.int64), F.cpu())
eid = g.edge_ids([0, 0, 0, 0, 0, 1, 2, 3, 4, 5],
[1, 2, 3, 4, 5, 9, 9, 9, 9, 9])
expected = F.asnumpy(expected)
eid = F.asnumpy(eid)
expected[eid] = 1
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
def test_multi_recv():
# basic recv test
g = generate_graph()
h = g.ndata['h']
g.register_message_func(message_func)
g.register_reduce_func(reduce_func)
g.register_apply_node_func(apply_node_func)
expected = F.copy_to(F.zeros((g.number_of_edges(),), dtype=F.int64), F.cpu())
# two separate round of send and recv
u = [4, 5, 6]
v = [9]
g.send((u, v))
eid = g.edge_ids(u, v)
expected = F.asnumpy(expected)
eid = F.asnumpy(eid)
expected[eid] = 1
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
g.recv(v)
expected[eid] = 0
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
u = [0]
v = [1, 2, 3]
g.send((u, v))
eid = g.edge_ids(u, v)
eid = F.asnumpy(eid)
expected[eid] = 1
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
g.recv(v)
expected[eid] = 0
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
h1 = g.ndata['h']
# one send, two recv
g.ndata['h'] = h
u = F.tensor([0, 0, 0, 4, 5, 6])
v = F.tensor([1, 2, 3, 9, 9, 9])
g.send((u, v))
eid = g.edge_ids(u, v)
eid = F.asnumpy(eid)
expected[eid] = 1
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
u = [4, 5, 6]
v = [9]
g.recv(v)
eid = g.edge_ids(u, v)
eid = F.asnumpy(eid)
expected[eid] = 0
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
u = [0]
v = [1, 2, 3]
g.recv(v)
eid = g.edge_ids(u, v)
eid = F.asnumpy(eid)
expected[eid] = 0
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
h2 = g.ndata['h']
assert F.allclose(h1, h2)
def test_multi_recv_0deg():
# test recv with 0deg nodes;
g = DGLGraph()
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype=dtype, ctx=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(_init2)
g.add_nodes(2)
g.add_edge(0, 1)
# recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata['h']
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full((5,), 4, F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# recv again on zero degree node
g.recv([0])
assert F.allclose(g.nodes[0].data['h'], F.full((5,), 8, F.float32))
# recv again on node with no incoming message
g.recv([1])
assert F.allclose(g.nodes[1].data['h'], F.sum(old, 0) * 4)
def test_send_twice_different_shape():
g = generate_graph()
def _message_1(edges):
return {'h': edges.src['h']}
def _message_2(edges):
return {'h': F.cat((edges.src['h'], edges.data['w']), dim=1)}
g.send(message_func=_message_1)
g.send(message_func=_message_2)
def test_send_twice_different_msg():
g = DGLGraph()
g.set_n_initializer(dgl.init.zero_initializer)
g.add_nodes(3)
g.add_edge(0, 1)
g.add_edge(2, 1)
def _message_a(edges):
return {'a': edges.src['a']}
def _message_b(edges):
return {'a': edges.src['a'] * 3}
def _reduce(nodes):
return {'a': F.max(nodes.mailbox['a'], 1)}
old_repr = F.randn((3, 5))
g.ndata['a'] = old_repr
g.send((0, 1), _message_a)
g.send((0, 1), _message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], old_repr[0] * 3)
g.ndata['a'] = old_repr
g.send((0, 1), _message_a)
g.send((2, 1), _message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], F.max(F.stack([old_repr[0], old_repr[2] * 3], 0), 0))
def test_send_twice_different_field():
g = DGLGraph()
g.set_n_initializer(dgl.init.zero_initializer)
g.add_nodes(2)
g.add_edge(0, 1)
def _message_a(edges):
return {'a': edges.src['a']}
def _message_b(edges):
return {'b': edges.src['b']}
def _reduce(nodes):
return {'a': F.sum(nodes.mailbox['a'], 1), 'b': F.sum(nodes.mailbox['b'], 1)}
old_a = F.randn((2, 5))
old_b = F.randn((2, 5))
g.set_n_repr({'a': old_a, 'b': old_b})
g.send((0, 1), _message_a)
g.send((0, 1), _message_b)
g.recv([1], _reduce)
new_repr = g.get_n_repr()
assert F.allclose(new_repr['a'][1], old_a[0])
assert F.allclose(new_repr['b'][1], old_b[0])
def test_dynamic_addition():
N = 3
D = 1
g = DGLGraph()
def _init(shape, dtype, ctx, ids):
return F.copy_to(F.astype(F.randn(shape), dtype), ctx)
g.set_n_initializer(_init)
g.set_e_initializer(_init)
def _message(edges):
return {'m' : edges.src['h1'] + edges.dst['h2'] + edges.data['h1'] +
edges.data['h2']}
def _reduce(nodes):
return {'h' : F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h' : nodes.data['h']}
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
# add nodes and edges
g.add_nodes(N)
g.ndata.update({'h1': F.randn((N, D)),
'h2': F.randn((N, D))})
g.add_nodes(3)
g.add_edge(0, 1)
g.add_edge(1, 0)
g.edata.update({'h1': F.randn((2, D)),
'h2': F.randn((2, D))})
g.send()
expected = F.copy_to(F.ones((g.number_of_edges(),), dtype=F.int64), F.cpu())
assert F.array_equal(g._get_msg_index().tousertensor(), expected)
# add more edges
g.add_edges([0, 2], [2, 0], {'h1': F.randn((2, D))})
g.send(([0, 2], [2, 0]))
g.recv(0)
g.add_edge(1, 2)
g.edges[4].data['h1'] = F.randn((1, D))
g.send((1, 2))
g.recv([1, 2])
h = g.ndata.pop('h')
# a complete round of send and recv
g.send()
g.recv()
assert F.allclose(h, g.ndata['h'])
def test_recv_no_send():
g = generate_graph()
g.recv(1, reduce_func)
# test recv after clear
g.clear()
g.add_nodes(3)
g.add_edges([0, 1], [1, 2])
g.set_n_initializer(dgl.init.zero_initializer)
g.ndata['h'] = F.randn((3, D))
g.send((1, 2), message_func)
expected = F.copy_to(F.zeros(2, dtype=F.int64), F.cpu())
expected = F.asnumpy(expected)
expected[1] = 1
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
g.recv(2, reduce_func)
expected[1] = 0
assert np.array_equal(g._get_msg_index().tonumpy(), expected)
def test_send_recv_after_conversion():
# test send and recv after converting from a graph with edges
g = generate_graph()
# nx graph
nxg = g.to_networkx(node_attrs=['h'])
g1 = DGLGraph()
# some random node and edges
g1.add_nodes(4)
g1.add_edges([1, 2], [2, 3])
g1.set_n_initializer(dgl.init.zero_initializer)
g1.from_networkx(nxg, node_attrs=['h'])
# sparse matrix
row, col= g.all_edges()
data = range(len(row))
n = g.number_of_nodes()
a = sp.coo_matrix(
(data, (F.zerocopy_to_numpy(row), F.zerocopy_to_numpy(col))),
shape=(n, n))
g2 = DGLGraph()
# some random node and edges
g2.add_nodes(5)
g2.add_edges([1, 2, 4], [2, 3, 0])
g2.set_n_initializer(dgl.init.zero_initializer)
g2.from_scipy_sparse_matrix(a)
g2.ndata['h'] = g.ndata['h']
# on dgl graph
g.send(message_func=message_func)
g.recv([0, 1, 3, 5], reduce_func=reduce_func,
apply_node_func=apply_node_func)
g.recv([0, 2, 4, 8], reduce_func=reduce_func,
apply_node_func=apply_node_func)
# nx
g1.send(message_func=message_func)
g1.recv([0, 1, 3, 5], reduce_func=reduce_func,
apply_node_func=apply_node_func)
g1.recv([0, 2, 4, 8], reduce_func=reduce_func,
apply_node_func=apply_node_func)
# sparse matrix
g2.send(message_func=message_func)
g2.recv([0, 1, 3, 5], reduce_func=reduce_func,
apply_node_func=apply_node_func)
g2.recv([0, 2, 4, 8], reduce_func=reduce_func,
apply_node_func=apply_node_func)
assert F.allclose(g.ndata['h'], g1.ndata['h'])
assert F.allclose(g.ndata['h'], g2.ndata['h'])
if __name__ == '__main__':
test_multi_send()
test_multi_recv()
test_multi_recv_0deg()
test_dynamic_addition()
test_send_twice_different_shape()
test_send_twice_different_msg()
test_send_twice_different_field()
test_recv_no_send()
test_send_recv_after_conversion()
tests/compute/test_nodeflow.py
View file @ 44089c8b
......@@ -28,40 +28,17 @@ def generate_rand_graph(n, connect_more=False, complete=False, add_self_loop=Fal
arr[0] = 1
arr[:, 0] = 1
if add_self_loop:
g = dgl.DGLGraph(arr, readonly=False)
g = dgl.DGLGraphStale(arr, readonly=False)
nodes = np.arange(g.number_of_nodes())
g.add_edges(nodes, nodes)
g.readonly()
else:
g = dgl.DGLGraph(arr, readonly=True)
g = dgl.DGLGraphStale(arr, readonly=True)
g.ndata['h1'] = F.randn((g.number_of_nodes(), 10))
g.edata['h2'] = F.randn((g.number_of_edges(), 3))
return g
def test_self_loop():
n = 100
num_hops = 2
g = generate_rand_graph(n, complete=True)
nf = create_mini_batch(g, num_hops, add_self_loop=True)
for i in range(1, nf.num_layers):
in_deg = nf.layer_in_degree(i)
deg = F.copy_to(F.ones(in_deg.shape, dtype=F.int64), F.cpu()) * n
assert_array_equal(F.asnumpy(in_deg), F.asnumpy(deg))
g = generate_rand_graph(n, complete=True, add_self_loop=True)
g = dgl.to_simple_graph(g)
nf = create_mini_batch(g, num_hops, add_self_loop=True)
for i in range(nf.num_blocks):
parent_eid = F.asnumpy(nf.block_parent_eid(i))
parent_nid = F.asnumpy(nf.layer_parent_nid(i + 1))
# The loop eid in the parent graph must exist in the block parent eid.
parent_loop_eid = F.asnumpy(g.edge_ids(parent_nid, parent_nid))
assert len(parent_loop_eid) == len(parent_nid)
for eid in parent_loop_eid:
assert eid in parent_eid
def create_mini_batch(g, num_hops, add_self_loop=False):
seed_ids = np.array([1, 2, 0, 3])
sampler = NeighborSampler(g, batch_size=4, expand_factor=g.number_of_nodes(),
......
tests/compute/test_pickle.py
View file @ 44089c8b
......@@ -9,7 +9,9 @@ import backend as F
import dgl.function as fn
import pickle
import io
import unittest
import unittest, pytest
import test_utils
from test_utils import parametrize_dtype, get_cases
def _assert_is_identical(g, g2):
assert g.is_readonly == g2.is_readonly
......@@ -147,130 +149,15 @@ def test_pickling_frame():
def _global_message_func(nodes):
return {'x': nodes.data['x']}
def test_pickling_graph():
# graph structures and frames are pickled
g = dgl.DGLGraph()
g.add_nodes(3)
src = F.tensor([0, 0])
dst = F.tensor([1, 2])
g.add_edges(src, dst)
x = F.randn((3, 7))
y = F.randn((3, 5))
a = F.randn((2, 6))
b = F.randn((2, 4))
g.ndata['x'] = x
g.ndata['y'] = y
g.edata['a'] = a
g.edata['b'] = b
# registered functions are pickled
g.register_message_func(_global_message_func)
reduce_func = fn.sum('x', 'x')
g.register_reduce_func(reduce_func)
# custom attributes should be pickled
g.foo = 2
new_g = _reconstruct_pickle(g)
_assert_is_identical(g, new_g)
assert new_g.foo == 2
assert new_g._message_func == _global_message_func
assert isinstance(new_g._reduce_func, type(reduce_func))
assert new_g._reduce_func._name == 'sum'
assert new_g._reduce_func.msg_field == 'x'
assert new_g._reduce_func.out_field == 'x'
# test batched graph with partial set case
g2 = dgl.DGLGraph()
g2.add_nodes(4)
src2 = F.tensor([0, 1])
dst2 = F.tensor([2, 3])
g2.add_edges(src2, dst2)
x2 = F.randn((4, 7))
y2 = F.randn((3, 5))
a2 = F.randn((2, 6))
b2 = F.randn((2, 4))
g2.ndata['x'] = x2
g2.nodes[[0, 1, 3]].data['y'] = y2
g2.edata['a'] = a2
g2.edata['b'] = b2
bg = dgl.batch([g, g2])
bg2 = _reconstruct_pickle(bg)
_assert_is_identical(bg, bg2)
new_g, new_g2 = dgl.unbatch(bg2)
_assert_is_identical(g, new_g)
_assert_is_identical(g2, new_g2)
# readonly graph
g = dgl.DGLGraph([(0, 1), (1, 2)], readonly=True)
new_g = _reconstruct_pickle(g)
_assert_is_identical(g, new_g)
# multigraph
g = dgl.DGLGraph([(0, 1), (0, 1), (1, 2)])
new_g = _reconstruct_pickle(g)
_assert_is_identical(g, new_g)
# readonly multigraph
g = dgl.DGLGraph([(0, 1), (0, 1), (1, 2)], readonly=True)
new_g = _reconstruct_pickle(g)
_assert_is_identical(g, new_g)
def test_pickling_nodeflow():
elist = [(0, 1), (1, 2), (2, 3), (3, 0)]
g = dgl.DGLGraph(elist, readonly=True)
g.ndata['x'] = F.randn((4, 5))
g.edata['y'] = F.randn((4, 3))
nf = contrib.sampling.sampler.create_full_nodeflow(g, 5)
nf.copy_from_parent() # add features
new_nf = _reconstruct_pickle(nf)
_assert_is_identical_nodeflow(nf, new_nf)
def test_pickling_batched_graph():
glist = [nx.path_graph(i + 5) for i in range(5)]
glist = [dgl.DGLGraph(g) for g in glist]
bg = dgl.batch(glist)
bg.ndata['x'] = F.randn((35, 5))
bg.edata['y'] = F.randn((60, 3))
new_bg = _reconstruct_pickle(bg)
_assert_is_identical_batchedgraph(bg, new_bg)
def test_pickling_heterograph():
# copied from test_heterograph.create_test_heterograph()
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
wishes_nx = nx.DiGraph()
wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
wishes_nx.add_edge('u0', 'g1', id=0)
wishes_nx.add_edge('u2', 'g0', id=1)
follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game')
wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game')
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')
g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
g.nodes['user'].data['u_h'] = F.randn((3, 4))
g.nodes['game'].data['g_h'] = F.randn((2, 5))
g.edges['plays'].data['p_h'] = F.randn((4, 6))
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
@pytest.mark.parametrize('g', get_cases(exclude=['dglgraph']))
def test_pickling_graph(g, idtype):
g = g.astype(idtype)
new_g = _reconstruct_pickle(g)
_assert_is_identical_hetero(g, new_g)
block = dgl.to_block(g, {'user': [1, 2], 'game': [0, 1], 'developer': []})
new_block = _reconstruct_pickle(block)
_assert_is_identical_hetero(block, new_block)
assert block.is_block
assert new_block.is_block
test_utils.check_graph_equal(g, new_g, check_feature=True)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_pickling_batched_heterograph():
# copied from test_heterograph.create_test_heterograph()
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
......@@ -296,8 +183,9 @@ def test_pickling_batched_heterograph():
bg = dgl.batch_hetero([g, g2])
new_bg = _reconstruct_pickle(bg)
_assert_is_identical_batchedhetero(bg, new_bg)
test_utils.check_graph_equal(bg, new_bg)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="Only test for pytorch format file")
def test_pickling_heterograph_index_compatibility():
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
......
tests/compute/test_propagate.py
View file @ 44089c8b
import dgl
import networkx as nx
import backend as F
import unittest
import utils as U
from utils import parametrize_dtype
def mfunc(edges):
return {'m' : edges.src['x']}
......@@ -10,18 +12,20 @@ def rfunc(nodes):
msg = F.sum(nodes.mailbox['m'], 1)
return {'x' : nodes.data['x'] + msg}
def test_prop_nodes_bfs():
g = dgl.DGLGraph(nx.path_graph(5))
g = dgl.graph(g.edges())
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_prop_nodes_bfs(idtype):
g = dgl.graph(nx.path_graph(5), idtype=idtype, device=F.ctx())
g.ndata['x'] = F.ones((5, 2))
dgl.prop_nodes_bfs(g, 0, message_func=mfunc, reduce_func=rfunc, apply_node_func=None)
# pull nodes using bfs order will result in a cumsum[i] + data[i] + data[i+1]
assert F.allclose(g.ndata['x'],
F.tensor([[2., 2.], [4., 4.], [6., 6.], [8., 8.], [9., 9.]]))
def test_prop_edges_dfs():
g = dgl.DGLGraph(nx.path_graph(5))
g = dgl.graph(g.edges())
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_prop_edges_dfs(idtype):
g = dgl.graph(nx.path_graph(5), idtype=idtype, device=F.ctx())
g.ndata['x'] = F.ones((5, 2))
dgl.prop_edges_dfs(g, 0, message_func=mfunc, reduce_func=rfunc, apply_node_func=None)
# snr using dfs results in a cumsum
......@@ -40,10 +44,11 @@ def test_prop_edges_dfs():
assert F.allclose(g.ndata['x'],
F.tensor([[3., 3.], [5., 5.], [7., 7.], [9., 9.], [5., 5.]]))
def test_prop_nodes_topo():
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_prop_nodes_topo(idtype):
# bi-directional chain
g = dgl.DGLGraph(nx.path_graph(5))
g = dgl.graph(g.edges())
g = dgl.graph(nx.path_graph(5), idtype=idtype, device=F.ctx())
assert U.check_fail(dgl.prop_nodes_topo, g) # has loop
# tree
......
tests/compute/test_randomwalk.py
View file @ 44089c8b
......@@ -12,7 +12,7 @@ def test_random_walk():
n_traces = 3
n_hops = 4
g = dgl.DGLGraph(edge_list, readonly=True)
g = dgl.DGLGraphStale(edge_list, readonly=True)
traces = dgl.contrib.sampling.random_walk(g, seeds, n_traces, n_hops)
traces = F.zerocopy_to_numpy(traces)
......@@ -31,7 +31,7 @@ def test_random_walk_with_restart():
seeds = [0, 1]
max_nodes = 10
g = dgl.DGLGraph(edge_list)
g = dgl.DGLGraphStale(edge_list)
# test normal RWR
traces = dgl.contrib.sampling.random_walk_with_restart(g, seeds, 0.2, max_nodes)
......@@ -61,9 +61,9 @@ def test_random_walk_with_restart():
assert (trace_diff % 2 == 0).all()
@parametrize_dtype
def test_metapath_random_walk(index_dtype):
g1 = dgl.bipartite(([0, 1, 2, 3], [0, 1, 2, 3]), 'a', 'ab', 'b', index_dtype=index_dtype)
g2 = dgl.bipartite(([0, 0, 1, 1, 2, 2, 3, 3], [1, 3, 2, 0, 3, 1, 0, 2]), 'b', 'ba', 'a', index_dtype=index_dtype)
def test_metapath_random_walk(idtype):
g1 = dgl.bipartite(([0, 1, 2, 3], [0, 1, 2, 3]), 'a', 'ab', 'b', idtype=idtype)
g2 = dgl.bipartite(([0, 0, 1, 1, 2, 2, 3, 3], [1, 3, 2, 0, 3, 1, 0, 2]), 'b', 'ba', 'a', idtype=idtype)
G = dgl.hetero_from_relations([g1, g2])
seeds = [0, 1]
traces = dgl.contrib.sampling.metapath_random_walk(G, ['ab', 'ba'] * 4, seeds, 3)
......
tests/compute/test_readout.py
View file @ 44089c8b
import dgl
import numpy as np
import backend as F
import networkx as nx
import unittest
def test_simple_readout():
g1 = dgl.DGLGraph()
g1.add_nodes(3)
g2 = dgl.DGLGraph()
g2.add_nodes(4) # no edges
g1.add_edges([0, 1, 2], [2, 0, 1])
n1 = F.randn((3, 5))
n2 = F.randn((4, 5))
e1 = F.randn((3, 5))
s1 = F.sum(n1, 0) # node sums
s2 = F.sum(n2, 0)
se1 = F.sum(e1, 0) # edge sums
m1 = F.mean(n1, 0) # node means
m2 = F.mean(n2, 0)
me1 = F.mean(e1, 0) # edge means
w1 = F.randn((3,))
w2 = F.randn((4,))
max1 = F.max(n1, 0)
max2 = F.max(n2, 0)
maxe1 = F.max(e1, 0)
ws1 = F.sum(n1 * F.unsqueeze(w1, 1), 0)
ws2 = F.sum(n2 * F.unsqueeze(w2, 1), 0)
wm1 = F.sum(n1 * F.unsqueeze(w1, 1), 0) / F.sum(F.unsqueeze(w1, 1), 0)
wm2 = F.sum(n2 * F.unsqueeze(w2, 1), 0) / F.sum(F.unsqueeze(w2, 1), 0)
g1.ndata['x'] = n1
g2.ndata['x'] = n2
g1.ndata['w'] = w1
g2.ndata['w'] = w2
g1.edata['x'] = e1
assert F.allclose(dgl.sum_nodes(g1, 'x'), s1)
assert F.allclose(dgl.sum_nodes(g1, 'x', 'w'), ws1)
assert F.allclose(dgl.sum_edges(g1, 'x'), se1)
assert F.allclose(dgl.mean_nodes(g1, 'x'), m1)
assert F.allclose(dgl.mean_nodes(g1, 'x', 'w'), wm1)
assert F.allclose(dgl.mean_edges(g1, 'x'), me1)
assert F.allclose(dgl.max_nodes(g1, 'x'), max1)
assert F.allclose(dgl.max_edges(g1, 'x'), maxe1)
g = dgl.batch([g1, g2])
s = dgl.sum_nodes(g, 'x')
m = dgl.mean_nodes(g, 'x')
max_bg = dgl.max_nodes(g, 'x')
assert F.allclose(s, F.stack([s1, s2], 0))
assert F.allclose(m, F.stack([m1, m2], 0))
assert F.allclose(max_bg, F.stack([max1, max2], 0))
ws = dgl.sum_nodes(g, 'x', 'w')
wm = dgl.mean_nodes(g, 'x', 'w')
assert F.allclose(ws, F.stack([ws1, ws2], 0))
assert F.allclose(wm, F.stack([wm1, wm2], 0))
s = dgl.sum_edges(g, 'x')
m = dgl.mean_edges(g, 'x')
max_bg_e = dgl.max_edges(g, 'x')
assert F.allclose(s, F.stack([se1, F.zeros(5)], 0))
assert F.allclose(m, F.stack([me1, F.zeros(5)], 0))
# TODO(zihao): fix -inf issue
# assert F.allclose(max_bg_e, F.stack([maxe1, F.zeros(5)], 0))
def test_topk_nodes():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(14))
feat0 = F.randn((g0.number_of_nodes(), 10))
g0.ndata['x'] = feat0
# to test the case where k > number of nodes.
dgl.topk_nodes(g0, 'x', 20, idx=-1)
# test correctness
val, indices = dgl.topk_nodes(g0, 'x', 5, idx=-1)
ground_truth = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 9, 10), 0, True)[:5], (5,))
assert F.allclose(ground_truth, indices)
g0.ndata.pop('x')
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(12))
feat1 = F.randn((g1.number_of_nodes(), 10))
bg = dgl.batch([g0, g1])
bg.ndata['x'] = F.cat([feat0, feat1], 0)
# to test the case where k > number of nodes.
dgl.topk_nodes(bg, 'x', 16, idx=1)
# test correctness
val, indices = dgl.topk_nodes(bg, 'x', 6, descending=False, idx=0)
ground_truth_0 = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 0, 1), 0, False)[:6], (6,))
ground_truth_1 = F.reshape(
F.argsort(F.slice_axis(feat1, -1, 0, 1), 0, False)[:6], (6,))
ground_truth = F.stack([ground_truth_0, ground_truth_1], 0)
assert F.allclose(ground_truth, indices)
# test idx=None
val, indices = dgl.topk_nodes(bg, 'x', 6, descending=True)
assert F.allclose(val, F.stack([F.topk(feat0, 6, 0), F.topk(feat1, 6, 0)], 0))
def test_topk_edges():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(14))
feat0 = F.randn((g0.number_of_edges(), 10))
g0.edata['x'] = feat0
# to test the case where k > number of edges.
dgl.topk_edges(g0, 'x', 30, idx=-1)
# test correctness
val, indices = dgl.topk_edges(g0, 'x', 7, idx=-1)
ground_truth = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 9, 10), 0, True)[:7], (7,))
assert F.allclose(ground_truth, indices)
g0.edata.pop('x')
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(12))
feat1 = F.randn((g1.number_of_edges(), 10))
bg = dgl.batch([g0, g1])
bg.edata['x'] = F.cat([feat0, feat1], 0)
# to test the case where k > number of edges.
dgl.topk_edges(bg, 'x', 33, idx=1)
# test correctness
val, indices = dgl.topk_edges(bg, 'x', 4, descending=False, idx=0)
ground_truth_0 = F.reshape(
F.argsort(F.slice_axis(feat0, -1, 0, 1), 0, False)[:4], (4,))
ground_truth_1 = F.reshape(
F.argsort(F.slice_axis(feat1, -1, 0, 1), 0, False)[:4], (4,))
ground_truth = F.stack([ground_truth_0, ground_truth_1], 0)
assert F.allclose(ground_truth, indices)
# test idx=None
val, indices = dgl.topk_edges(bg, 'x', 6, descending=True)
assert F.allclose(val, F.stack([F.topk(feat0, 6, 0), F.topk(feat1, 6, 0)], 0))
def test_softmax_nodes():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(9))
feat0 = F.randn((g0.number_of_nodes(), 10))
g0.ndata['x'] = feat0
ground_truth = F.softmax(feat0, dim=0)
assert F.allclose(dgl.softmax_nodes(g0, 'x'), ground_truth)
g0.ndata.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_nodes(), 10))
feat2 = F.randn((g2.number_of_nodes(), 10))
feat4 = F.randn((g4.number_of_nodes(), 10))
bg.ndata['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_nodes(bg, 'x'), ground_truth)
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_broadcast_nodes():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(10))
feat0 = F.randn((1, 40))
ground_truth = F.stack([feat0] * g0.number_of_nodes(), 0)
assert F.allclose(dgl.broadcast_nodes(g0, feat0), ground_truth)
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(3))
g2 = dgl.DGLGraph()
g3 = dgl.DGLGraph(nx.path_graph(12))
bg = dgl.batch([g0, g1, g2, g3])
feat1 = F.randn((1, 40))
feat2 = F.randn((1, 40))
feat3 = F.randn((1, 40))
ground_truth = F.cat(
[feat0] * g0.number_of_nodes() +\
[feat1] * g1.number_of_nodes() +\
[feat2] * g2.number_of_nodes() +\
[feat3] * g3.number_of_nodes(), 0
)
assert F.allclose(dgl.broadcast_nodes(
bg, F.cat([feat0, feat1, feat2, feat3], 0)
), ground_truth)
def test_broadcast_edges():
# test#1: basic
g0 = dgl.DGLGraph(nx.path_graph(10))
feat0 = F.randn((1, 40))
ground_truth = F.stack([feat0] * g0.number_of_edges(), 0)
assert F.allclose(dgl.broadcast_edges(g0, feat0), ground_truth)
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(3))
g2 = dgl.DGLGraph()
g3 = dgl.DGLGraph(nx.path_graph(12))
bg = dgl.batch([g0, g1, g2, g3])
feat1 = F.randn((1, 40))
feat2 = F.randn((1, 40))
feat3 = F.randn((1, 40))
ground_truth = F.cat(
[feat0] * g0.number_of_edges() +\
[feat1] * g1.number_of_edges() +\
[feat2] * g2.number_of_edges() +\
[feat3] * g3.number_of_edges(), 0
)
assert F.allclose(dgl.broadcast_edges(
bg, F.cat([feat0, feat1, feat2, feat3], 0)
), ground_truth)
if __name__ == '__main__':
test_simple_readout()
test_topk_nodes()
test_topk_edges()
test_softmax_nodes()
test_softmax_edges()
test_broadcast_nodes()
test_broadcast_edges()
import pytest
from test_utils.graph_cases import get_cases
from utils import parametrize_dtype
@parametrize_dtype
def test_sum_case1(idtype):
# NOTE: If you want to update this test case, remember to update the docstring
# example too!!!
g1 = dgl.graph(([0, 1], [1, 0]), idtype=idtype, device=F.ctx())
g1.ndata['h'] = F.tensor([1., 2.])
g2 = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g2.ndata['h'] = F.tensor([1., 2., 3.])
bg = dgl.batch([g1, g2])
bg.ndata['w'] = F.tensor([.1, .2, .1, .5, .2])
assert F.allclose(F.tensor([3.]), dgl.sum_nodes(g1, 'h'))
assert F.allclose(F.tensor([3., 6.]), dgl.sum_nodes(bg, 'h'))
assert F.allclose(F.tensor([.5, 1.7]), dgl.sum_nodes(bg, 'h', 'w'))
@parametrize_dtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph']))
@pytest.mark.parametrize('reducer', ['sum', 'max', 'mean'])
def test_reduce_readout(g, idtype, reducer):
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.readout_nodes(g, 'h', op=reducer)
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = dgl.readout_nodes(sg, 'h', op=reducer)
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
x = getattr(dgl, '{}_nodes'.format(reducer))(g, 'h')
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = getattr(dgl, '{}_nodes'.format(reducer))(sg, 'h')
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
# Test.2: edge readout
x = dgl.readout_edges(g, 'h', op=reducer)
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = dgl.readout_edges(sg, 'h', op=reducer)
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
x = getattr(dgl, '{}_edges'.format(reducer))(g, 'h')
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = getattr(dgl, '{}_edges'.format(reducer))(sg, 'h')
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
@parametrize_dtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph']))
@pytest.mark.parametrize('reducer', ['sum', 'max', 'mean'])
def test_weighted_reduce_readout(g, idtype, reducer):
g = g.astype(idtype).to(F.ctx())
g.ndata['h'] = F.randn((g.number_of_nodes(), 3))
g.ndata['w'] = F.randn((g.number_of_nodes(), 1))
g.edata['h'] = F.randn((g.number_of_edges(), 2))
g.edata['w'] = F.randn((g.number_of_edges(), 1))
# Test.1: node readout
x = dgl.readout_nodes(g, 'h', 'w', op=reducer)
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = dgl.readout_nodes(sg, 'h', 'w', op=reducer)
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
x = getattr(dgl, '{}_nodes'.format(reducer))(g, 'h', 'w')
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = getattr(dgl, '{}_nodes'.format(reducer))(sg, 'h', 'w')
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
# Test.2: edge readout
x = dgl.readout_edges(g, 'h', 'w', op=reducer)
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = dgl.readout_edges(sg, 'h', 'w', op=reducer)
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
x = getattr(dgl, '{}_edges'.format(reducer))(g, 'h', 'w')
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = getattr(dgl, '{}_edges'.format(reducer))(sg, 'h', 'w')
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
@parametrize_dtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph']))
@pytest.mark.parametrize('descending', [True, False])
def test_topk(g, idtype, descending):
g = g.astype(idtype).to(F.ctx())
g.ndata['x'] = F.randn((g.number_of_nodes(), 3))
# Test.1: to test the case where k > number of nodes.
dgl.topk_nodes(g, 'x', 100, sortby=-1)
# Test.2: test correctness
min_nnodes = F.asnumpy(g.batch_num_nodes()).min()
if min_nnodes <= 1:
return
k = min_nnodes - 1
val, indices = dgl.topk_nodes(g, 'x', k, descending=descending, sortby=-1)
print(k)
print(g.ndata['x'])
print('val', val)
print('indices', indices)
subg = dgl.unbatch(g)
subval, subidx = [], []
for sg in subg:
subx = F.asnumpy(sg.ndata['x'])
ai = np.argsort(subx[:,-1:].flatten())
if descending:
ai = np.ascontiguousarray(ai[::-1])
subx = np.expand_dims(subx[ai[:k]], 0)
subval.append(F.tensor(subx))
subidx.append(F.tensor(np.expand_dims(ai[:k], 0)))
print(F.cat(subval, dim=0))
assert F.allclose(val, F.cat(subval, dim=0))
assert F.allclose(indices, F.cat(subidx, dim=0))
# Test.3: sorby=None
dgl.topk_nodes(g, 'x', k, sortby=None)
g.edata['x'] = F.randn((g.number_of_edges(), 3))
# Test.4: topk edges where k > number of edges.
dgl.topk_edges(g, 'x', 100, sortby=-1)
# Test.5: topk edges test correctness
min_nedges = F.asnumpy(g.batch_num_edges()).min()
if min_nedges <= 1:
return
k = min_nedges - 1
val, indices = dgl.topk_edges(g, 'x', k, descending=descending, sortby=-1)
print(k)
print(g.edata['x'])
print('val', val)
print('indices', indices)
subg = dgl.unbatch(g)
subval, subidx = [], []
for sg in subg:
subx = F.asnumpy(sg.edata['x'])
ai = np.argsort(subx[:,-1:].flatten())
if descending:
ai = np.ascontiguousarray(ai[::-1])
subx = np.expand_dims(subx[ai[:k]], 0)
subval.append(F.tensor(subx))
subidx.append(F.tensor(np.expand_dims(ai[:k], 0)))
print(F.cat(subval, dim=0))
assert F.allclose(val, F.cat(subval, dim=0))
assert F.allclose(indices, F.cat(subidx, dim=0))
@parametrize_dtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph']))
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))
@parametrize_dtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph']))
def test_broadcast(idtype, g):
g = g.astype(idtype).to(F.ctx())
gfeat = F.randn((g.batch_size, 3))
# Test.0: broadcast_nodes
g.ndata['h'] = dgl.broadcast_nodes(g, gfeat)
subg = dgl.unbatch(g)
for i, sg in enumerate(subg):
assert F.allclose(sg.ndata['h'],
F.repeat(F.reshape(gfeat[i], (1,3)), sg.number_of_nodes(), dim=0))
# Test.1: broadcast_edges
g.edata['h'] = dgl.broadcast_edges(g, gfeat)
subg = dgl.unbatch(g)
for i, sg in enumerate(subg):
assert F.allclose(sg.edata['h'],
F.repeat(F.reshape(gfeat[i], (1,3)), sg.number_of_edges(), dim=0))
tests/compute/test_removal.py
View file @ 44089c8b
......@@ -3,9 +3,12 @@ import backend as F
import networkx as nx
import numpy as np
import dgl
from test_utils import parametrize_dtype
def test_node_removal():
@parametrize_dtype
def test_node_removal(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
g.add_edge(0, 0)
assert g.number_of_nodes() == 10
......@@ -26,8 +29,10 @@ def test_node_removal():
assert g.number_of_nodes() == 7
assert F.array_equal(g.ndata['id'], F.tensor([0, 7, 8, 9, 0, 0, 0]))
def test_multigraph_node_removal():
@parametrize_dtype
def test_multigraph_node_removal(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(5)
for i in range(5):
g.add_edge(i, i)
......@@ -52,8 +57,10 @@ def test_multigraph_node_removal():
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 6
def test_multigraph_edge_removal():
@parametrize_dtype
def test_multigraph_edge_removal(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(5)
for i in range(5):
g.add_edge(i, i)
......@@ -77,8 +84,10 @@ def test_multigraph_edge_removal():
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 8
def test_edge_removal():
@parametrize_dtype
def test_edge_removal(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(5)
for i in range(5):
for j in range(5):
......@@ -103,8 +112,10 @@ def test_edge_removal():
assert g.number_of_edges() == 11
assert F.array_equal(g.edata['id'], F.tensor([0, 1, 10, 11, 12, 20, 21, 22, 23, 24, 0]))
def test_node_and_edge_removal():
@parametrize_dtype
def test_node_and_edge_removal(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
for i in range(10):
for j in range(10):
......@@ -140,46 +151,54 @@ def test_node_and_edge_removal():
assert g.number_of_nodes() == 10
assert g.number_of_edges() == 48
def test_node_frame():
@parametrize_dtype
def test_node_frame(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
data = np.random.rand(10, 3)
new_data = data.take([0, 1, 2, 7, 8, 9], axis=0)
g.ndata['h'] = F.zerocopy_from_numpy(data)
g.ndata['h'] = F.tensor(data)
# remove nodes
g.remove_nodes(range(3, 7))
assert F.allclose(g.ndata['h'], F.zerocopy_from_numpy(new_data))
assert F.allclose(g.ndata['h'], F.tensor(new_data))
def test_edge_frame():
@parametrize_dtype
def test_edge_frame(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
g.add_edges(list(range(10)), list(range(1, 10)) + [0])
data = np.random.rand(10, 3)
new_data = data.take([0, 1, 2, 7, 8, 9], axis=0)
g.edata['h'] = F.zerocopy_from_numpy(data)
g.edata['h'] = F.tensor(data)
# remove edges
g.remove_edges(range(3, 7))
assert F.allclose(g.edata['h'], F.zerocopy_from_numpy(new_data))
assert F.allclose(g.edata['h'], F.tensor(new_data))
def test_frame_size():
@parametrize_dtype
def test_issue1287(idtype):
# reproduce https://github.com/dmlc/dgl/issues/1287.
# remove nodes
# setting features after remove nodes
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(5)
g.add_edges([0, 2, 3, 1, 1], [1, 0, 3, 1, 0])
g.remove_nodes([0, 1])
assert g._node_frame.num_rows == 3
assert g._edge_frame.num_rows == 1
g.ndata['h'] = F.randn((g.number_of_nodes(), 3))
g.edata['h'] = F.randn((g.number_of_edges(), 2))
# remove edges
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(5)
g.add_edges([0, 2, 3, 1, 1], [1, 0, 3, 1, 0])
g.remove_edges([0, 1])
assert g._node_frame.num_rows == 5
assert g._edge_frame.num_rows == 3
g = g.to(F.ctx())
g.ndata['h'] = F.randn((g.number_of_nodes(), 3))
g.edata['h'] = F.randn((g.number_of_edges(), 2))
if __name__ == '__main__':
test_node_removal()
......
tests/compute/test_sampler.py
View file @ 44089c8b
......@@ -10,7 +10,7 @@ np.random.seed(42)
def generate_rand_graph(n):
arr = (sp.sparse.random(n, n, density=0.1, format='coo') != 0).astype(np.int64)
return dgl.DGLGraph(arr, readonly=True)
return dgl.DGLGraphStale(arr, readonly=True)
def test_create_full():
g = generate_rand_graph(100)
......@@ -171,7 +171,7 @@ def test_nonuniform_neighbor_sampler():
if edge not in edges:
edges.append(edge)
src, dst = zip(*edges)
g = dgl.DGLGraph()
g = dgl.DGLGraphStale()
g.add_nodes(100)
g.add_edges(src, dst)
g.readonly()
......
tests/compute/test_sampling.py
View file @ 44089c8b
......@@ -17,7 +17,7 @@ def check_random_walk(g, metapath, traces, ntypes, prob=None):
traces[i, j], traces[i, j+1], etype=metapath[j])
if prob is not None and prob in g.edges[metapath[j]].data:
p = F.asnumpy(g.edges[metapath[j]].data['p'])
eids = g.edge_id(traces[i, j], traces[i, j+1], etype=metapath[j])
eids = g.edge_ids(traces[i, j], traces[i, j+1], etype=metapath[j])
assert p[eids] != 0
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU random walk not implemented")
......@@ -113,6 +113,7 @@ def test_pack_traces():
assert F.array_equal(result[2], F.tensor([2, 7], dtype=F.int64))
assert F.array_equal(result[3], F.tensor([0, 2], dtype=F.int64))
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_pinsage_sampling():
def _test_sampler(g, sampler, ntype):
neighbor_g = sampler(F.tensor([0, 2], dtype=F.int64))
......@@ -228,7 +229,7 @@ def _test_sample_neighbors(hypersparse):
assert subg.number_of_edges() == 4
u, v = subg.edges()
assert set(F.asnumpy(F.unique(v))) == {0, 1}
assert F.array_equal(g.has_edges_between(u, v), F.ones((4,), dtype=F.int64))
assert F.array_equal(F.astype(g.has_edges_between(u, v), F.int64), F.ones((4,), dtype=F.int64))
assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
if not replace:
......@@ -258,7 +259,7 @@ def _test_sample_neighbors(hypersparse):
assert subg.number_of_edges() == num_edges
u, v = subg.edges()
assert set(F.asnumpy(F.unique(v))) == {0, 2}
assert F.array_equal(g.has_edges_between(u, v), F.ones((num_edges,), dtype=F.int64))
assert F.array_equal(F.astype(g.has_edges_between(u, v), F.int64), F.ones((num_edges,), dtype=F.int64))
assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
if not replace:
......@@ -326,7 +327,7 @@ def _test_sample_neighbors_outedge(hypersparse):
assert subg.number_of_edges() == 4
u, v = subg.edges()
assert set(F.asnumpy(F.unique(u))) == {0, 1}
assert F.array_equal(g.has_edges_between(u, v), F.ones((4,), dtype=F.int64))
assert F.array_equal(F.astype(g.has_edges_between(u, v), F.int64), F.ones((4,), dtype=F.int64))
assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
if not replace:
......@@ -356,7 +357,7 @@ def _test_sample_neighbors_outedge(hypersparse):
assert subg.number_of_edges() == num_edges
u, v = subg.edges()
assert set(F.asnumpy(F.unique(u))) == {0, 2}
assert F.array_equal(g.has_edges_between(u, v), F.ones((num_edges,), dtype=F.int64))
assert F.array_equal(F.astype(g.has_edges_between(u, v), F.int64), F.ones((num_edges,), dtype=F.int64))
assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
if not replace:
......
tests/compute/test_serialize.py
View file @ 44089c8b
......@@ -5,6 +5,7 @@ import time
import tempfile
import os
import pytest
import unittest
from dgl import DGLGraph
import dgl
......@@ -34,6 +35,7 @@ def construct_graph(n, is_hetero):
return g_list
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@pytest.mark.parametrize('is_hetero', [True, False])
def test_graph_serialize_with_feature(is_hetero):
num_graphs = 100
......@@ -75,6 +77,7 @@ def test_graph_serialize_with_feature(is_hetero):
os.unlink(path)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@pytest.mark.parametrize('is_hetero', [True, False])
def test_graph_serialize_without_feature(is_hetero):
num_graphs = 100
......@@ -102,6 +105,7 @@ def test_graph_serialize_without_feature(is_hetero):
os.unlink(path)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@pytest.mark.parametrize('is_hetero', [True, False])
def test_graph_serialize_with_labels(is_hetero):
num_graphs = 100
......@@ -212,10 +216,10 @@ def test_load_old_files2():
assert np.allclose(F.asnumpy(load_edges[1]), edges1)
def create_heterographs(index_dtype):
def create_heterographs(idtype):
g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user',
'follows', index_dtype=index_dtype, restrict_format='any')
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', index_dtype=index_dtype, restrict_format='csr')
'follows', idtype=idtype)
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', idtype=idtype).formats('csr')
g_x.nodes['user'].data['h'] = F.randn((4, 3))
g_x.edges['follows'].data['w'] = F.randn((3, 2))
g_y.nodes['user'].data['hh'] = F.ones((4, 5))
......@@ -223,11 +227,11 @@ def create_heterographs(index_dtype):
g = dgl.hetero_from_relations([g_x, g_y])
return [g, g_x, g_y]
def create_heterographs2(index_dtype):
def create_heterographs2(idtype):
g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user',
'follows', index_dtype=index_dtype, restrict_format='any')
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', index_dtype=index_dtype, restrict_format='csr')
g_z = dgl.bipartite(([0, 1, 3], [2, 3, 4]), 'user', 'knows', 'knowledge', index_dtype=index_dtype)
'follows', idtype=idtype)
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', idtype=idtype).formats('csr')
g_z = dgl.bipartite(([0, 1, 3], [2, 3, 4]), 'user', 'knows', 'knowledge', idtype=idtype)
g_x.nodes['user'].data['h'] = F.randn((4, 3))
g_x.edges['follows'].data['w'] = F.randn((3, 2))
g_y.nodes['user'].data['hh'] = F.ones((4, 5))
......@@ -253,16 +257,17 @@ def test_deserialize_old_heterograph_file():
def create_old_heterograph_files():
path = os.path.join(
os.path.dirname(__file__), "data/hetero1.bin")
g_list0 = create_heterographs("int64") + create_heterographs("int32")
g_list0 = create_heterographs(F.int64) + create_heterographs(F.int32)
labels_dict = {"graph_label": F.ones(54)}
save_graphs(path, g_list0, labels_dict)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_serialize_heterograph():
f = tempfile.NamedTemporaryFile(delete=False)
path = f.name
f.close()
g_list0 = create_heterographs2("int64") + create_heterographs2("int32")
g_list0 = create_heterographs2(F.int64) + create_heterographs2(F.int32)
save_graphs(path, g_list0)
g_list, _ = load_graphs(path)
......@@ -271,8 +276,9 @@ def test_serialize_heterograph():
for i in range(len(g_list0)):
for j, etypes in enumerate(g_list0[i].canonical_etypes):
assert g_list[i].canonical_etypes[j] == etypes
assert g_list[1].restrict_format() == 'any'
assert g_list[2].restrict_format() == 'csr'
#assert g_list[1].restrict_format() == 'any'
#assert g_list[2].restrict_format() == 'csr'
assert g_list[4].idtype == F.int32
assert np.allclose(
F.asnumpy(g_list[2].nodes['user'].data['hh']), np.ones((4, 5)))
......@@ -291,15 +297,16 @@ def test_serialize_heterograph():
os.unlink(path)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@pytest.mark.skip(reason="lack of permission on CI")
def test_serialize_heterograph_s3():
path = "s3://dglci-data-test/graph2.bin"
g_list0 = create_heterographs("int64") + create_heterographs("int32")
g_list0 = create_heterographs(F.int64) + create_heterographs(F.int32)
save_graphs(path, g_list0)
g_list = load_graphs(path, [0, 2, 5])
assert g_list[0].idtype == F.int64
assert g_list[1].restrict_format() == 'csr'
#assert g_list[1].restrict_format() == 'csr'
assert np.allclose(
F.asnumpy(g_list[1].nodes['user'].data['hh']), np.ones((4, 5)))
assert np.allclose(
......
tests/compute/test_sparse.py
View file @ 44089c8b
......@@ -5,7 +5,8 @@ import random
import pytest
import networkx as nx
import backend as F
import numpy as np
import numpy as np
from utils import parametrize_dtype
random.seed(42)
np.random.seed(42)
......@@ -98,13 +99,10 @@ sddmm_shapes = [
@pytest.mark.parametrize('msg', ['add', 'sub', 'mul', 'div', 'copy_lhs', 'copy_rhs'])
@pytest.mark.parametrize('reducer', ['sum', 'min', 'max'])
@parametrize_dtype
def test_spmm(g, shp, msg, reducer, index_dtype):
if dgl.backend.backend_name == 'tensorflow' and (reducer in ['min', 'max'] or index_dtype == 'int32'):
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.
if index_dtype == 'int32':
g = g.int()
else:
g = g.long()
print(g)
print(g.idtype)
......@@ -164,15 +162,12 @@ def test_spmm(g, shp, msg, reducer, index_dtype):
@pytest.mark.parametrize('rhs_target', ['u', 'v', 'e'])
@pytest.mark.parametrize('msg', ['add', 'sub', 'mul', 'div', 'dot', 'copy_lhs', 'copy_rhs'])
@parametrize_dtype
def test_sddmm(g, shp, lhs_target, rhs_target, msg, index_dtype):
def test_sddmm(g, shp, lhs_target, rhs_target, msg, idtype):
if lhs_target == rhs_target:
return
g = g.astype(idtype).to(F.ctx())
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)
......@@ -234,4 +229,4 @@ def test_sddmm(g, shp, lhs_target, rhs_target, msg, index_dtype):
if 'm' in g.edata: g.edata.pop('m')
if __name__ == '__main__':
test_spmm(graphs[0], spmm_shapes[5], 'copy_lhs', 'sum')
test_spmm(F.int32, graphs[0], spmm_shapes[5], 'copy_lhs', 'sum')
tests/compute/test_specialization.py
View file @ 44089c8b
......@@ -3,11 +3,13 @@ import scipy.sparse as sp
import dgl
import dgl.function as fn
import backend as F
from test_utils import parametrize_dtype
D = 5
def generate_graph():
def generate_graph(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
# create a graph where 0 is the source and 9 is the sink
for i in range(1, 9):
......@@ -15,12 +17,13 @@ def generate_graph():
g.add_edge(i, 9)
# add a back flow from 9 to 0
g.add_edge(9, 0)
g.set_n_repr({'f1' : F.randn((10,)), 'f2' : F.randn((10, D))})
g.ndata.update({'f1' : F.randn((10,)), 'f2' : F.randn((10, D))})
weights = F.randn((17,))
g.set_e_repr({'e1': weights, 'e2': F.unsqueeze(weights, 1)})
g.edata.update({'e1': weights, 'e2': F.unsqueeze(weights, 1)})
return g
def test_v2v_update_all():
@parametrize_dtype
def test_v2v_update_all(idtype):
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
......@@ -36,12 +39,12 @@ def test_v2v_update_all():
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = generate_graph()
g = generate_graph(idtype)
# update all
v1 = g.ndata[fld]
g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.ndata.update({fld : v1})
g.update_all(message_func, reduce_func, apply_func)
v3 = g.ndata[fld]
assert F.allclose(v2, v3)
......@@ -50,7 +53,7 @@ def test_v2v_update_all():
g.update_all(fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.ndata.update({fld : v1})
g.update_all(message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v4)
......@@ -59,9 +62,10 @@ def test_v2v_update_all():
# test 2d node features
_test('f2')
def test_v2v_snr():
u = F.tensor([0, 0, 0, 3, 4, 9])
v = F.tensor([1, 2, 3, 9, 9, 0])
@parametrize_dtype
def test_v2v_snr(idtype):
u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
......@@ -77,13 +81,13 @@ def test_v2v_snr():
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = generate_graph()
g = generate_graph(idtype)
# send and recv
v1 = g.ndata[fld]
g.send_and_recv((u, v), fn.copy_src(src=fld, out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.ndata.update({fld : v1})
g.send_and_recv((u, v), message_func, reduce_func, apply_func)
v3 = g.ndata[fld]
assert F.allclose(v2, v3)
......@@ -92,7 +96,7 @@ def test_v2v_snr():
g.send_and_recv((u, v), fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.ndata.update({fld : v1})
g.send_and_recv((u, v), message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v4)
......@@ -102,8 +106,9 @@ def test_v2v_snr():
_test('f2')
def test_v2v_pull():
nodes = F.tensor([1, 2, 3, 9])
@parametrize_dtype
def test_v2v_pull(idtype):
nodes = F.tensor([1, 2, 3, 9], idtype)
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
......@@ -119,7 +124,7 @@ def test_v2v_pull():
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = generate_graph()
g = generate_graph(idtype)
# send and recv
v1 = g.ndata[fld]
g.pull(nodes, fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
......@@ -142,7 +147,8 @@ def test_v2v_pull():
# test 2d node features
_test('f2')
def test_v2v_update_all_multi_fn():
@parametrize_dtype
def test_v2v_update_all_multi_fn(idtype):
def message_func(edges):
return {'m2': edges.src['f2']}
......@@ -152,8 +158,8 @@ def test_v2v_update_all_multi_fn():
def reduce_func(nodes):
return {'v1': F.sum(nodes.mailbox['m2'], 1)}
g = generate_graph()
g.set_n_repr({'v1' : F.zeros((10,)), 'v2' : F.zeros((10,))})
g = generate_graph(idtype)
g.ndata.update({'v1' : F.zeros((10,)), 'v2' : F.zeros((10,))})
fld = 'f2'
g.update_all(message_func, reduce_func)
......@@ -181,9 +187,10 @@ def test_v2v_update_all_multi_fn():
v2 = g.ndata['v2']
assert F.allclose(v1, v2)
def test_v2v_snr_multi_fn():
u = F.tensor([0, 0, 0, 3, 4, 9])
v = F.tensor([1, 2, 3, 9, 9, 0])
@parametrize_dtype
def test_v2v_snr_multi_fn(idtype):
u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
def message_func(edges):
return {'m2': edges.src['f2']}
......@@ -194,8 +201,8 @@ def test_v2v_snr_multi_fn():
def reduce_func(nodes):
return {'v1' : F.sum(nodes.mailbox['m2'], 1)}
g = generate_graph()
g.set_n_repr({'v1' : F.zeros((10, D)), 'v2' : F.zeros((10, D)),
g = generate_graph(idtype)
g.ndata.update({'v1' : F.zeros((10, D)), 'v2' : F.zeros((10, D)),
'v3' : F.zeros((10, D))})
fld = 'f2'
......@@ -229,7 +236,8 @@ def test_v2v_snr_multi_fn():
v2 = g.ndata['v2']
assert F.allclose(v1, v2)
def test_e2v_update_all_multi_fn():
@parametrize_dtype
def test_e2v_update_all_multi_fn(idtype):
def _test(fld):
def message_func(edges):
return {'m1' : edges.src[fld] + edges.dst[fld],
......@@ -244,19 +252,19 @@ def test_e2v_update_all_multi_fn():
def apply_func_2(nodes):
return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
g = generate_graph()
g = generate_graph(idtype)
# update all
v1 = g.get_n_repr()[fld]
v1 = g.ndata[fld]
# no specialization
g.update_all(message_func, reduce_func, apply_func)
v2 = g.get_n_repr()[fld]
v2 = g.ndata[fld]
# user break reduce func into 2 builtin
g.set_n_repr({fld : v1})
g.ndata.update({fld : v1})
g.update_all(message_func,
[fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
apply_func_2)
v3 = g.get_n_repr()[fld]
v3 = g.ndata[fld]
assert F.allclose(v2, v3)
......@@ -265,9 +273,10 @@ def test_e2v_update_all_multi_fn():
# test 2d node features
_test('f2')
def test_e2v_snr_multi_fn():
u = F.tensor([0, 0, 0, 3, 4, 9])
v = F.tensor([1, 2, 3, 9, 9, 0])
@parametrize_dtype
def test_e2v_snr_multi_fn(idtype):
u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
def _test(fld):
def message_func(edges):
return {'m1' : edges.src[fld] + edges.dst[fld],
......@@ -282,59 +291,19 @@ def test_e2v_snr_multi_fn():
def apply_func_2(nodes):
return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
g = generate_graph()
g = generate_graph(idtype)
# send_and_recv
v1 = g.get_n_repr()[fld]
v1 = g.ndata[fld]
# no specialization
g.send_and_recv((u, v), message_func, reduce_func, apply_func)
v2 = g.get_n_repr()[fld]
v2 = g.ndata[fld]
# user break reduce func into 2 builtin
g.set_n_repr({fld : v1})
g.ndata.update({fld : v1})
g.send_and_recv((u, v), message_func,
[fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
apply_func_2)
v3 = g.get_n_repr()[fld]
assert F.allclose(v2, v3)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
def test_e2v_recv_multi_fn():
u = F.tensor([0, 0, 0, 3, 4, 9])
v = F.tensor([1, 2, 3, 9, 9, 0])
def _test(fld):
def message_func(edges):
return {'m1' : edges.src[fld] + edges.dst[fld],
'm2' : edges.src[fld] * edges.dst[fld]}
def reduce_func(nodes):
return {fld : F.sum(nodes.mailbox['m1'] + nodes.mailbox['m2'], 1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
def apply_func_2(nodes):
return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
g = generate_graph()
# recv
v1 = g.get_n_repr()[fld]
# no specialization
g.send((u, v), message_func)
g.recv([0,1,2,3,9], reduce_func, apply_func)
v2 = g.get_n_repr()[fld]
# user break reduce func into 2 builtin
g.set_n_repr({fld : v1})
g.send((u, v), message_func)
g.recv([0,1,2,3,9],
[fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
apply_func_2)
v3 = g.get_n_repr()[fld]
v3 = g.ndata[fld]
assert F.allclose(v2, v3)
......@@ -343,9 +312,11 @@ def test_e2v_recv_multi_fn():
# test 2d node features
_test('f2')
def test_update_all_multi_fallback():
@parametrize_dtype
def test_update_all_multi_fallback(idtype):
# create a graph with zero in degree nodes
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
for i in range(1, 9):
g.add_edge(0, i)
......@@ -399,9 +370,11 @@ def test_update_all_multi_fallback():
assert F.allclose(o1, g.ndata.pop('o1'))
assert F.allclose(o2, g.ndata.pop('o2'))
def test_pull_multi_fallback():
@parametrize_dtype
def test_pull_multi_fallback(idtype):
# create a graph with zero in degree nodes
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
for i in range(1, 9):
g.add_edge(0, i)
......@@ -465,7 +438,8 @@ def test_pull_multi_fallback():
nodes = [0, 1, 2, 9]
_pull_nodes(nodes)
def test_spmv_3d_feat():
@parametrize_dtype
def test_spmv_3d_feat(idtype):
def src_mul_edge_udf(edges):
return {'sum': edges.src['h'] * F.unsqueeze(F.unsqueeze(edges.data['h'], 1), 1)}
......@@ -476,6 +450,7 @@ def test_spmv_3d_feat():
p = 0.1
a = sp.random(n, n, p, data_rvs=lambda n: np.ones(n))
g = dgl.DGLGraph(a)
g = g.astype(idtype).to(F.ctx())
m = g.number_of_edges()
# test#1: v2v with adj data
......
tests/compute/test_to_device.py deleted 100644 → 0
View file @ 015acfd2
import dgl
import backend as F
import unittest
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
def test_to_device():
g = dgl.DGLGraph()
g.add_nodes(5, {'h' : F.ones((5, 2))})
g.add_edges([0, 1], [1, 2], {'m' : F.ones((2, 2))})
if F.is_cuda_available():
g = g.to(F.cuda())
assert g is not None
if __name__ == '__main__':
test_to_device()
tests/compute/test_transform.py
View file @ 44089c8b
from scipy import sparse as spsp
import unittest
import networkx as nx
import numpy as np
import dgl
......@@ -9,37 +8,28 @@ from dgl.graph_index import from_scipy_sparse_matrix
import unittest
from utils import parametrize_dtype
from test_heterograph import create_test_heterograph4, create_test_heterograph5, create_test_heterograph6
D = 5
# line graph related
def test_line_graph():
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_line_graph1():
N = 5
G = dgl.DGLGraph(nx.star_graph(N))
G.edata['h'] = F.randn((2 * N, D))
n_edges = G.number_of_edges()
L = G.line_graph(shared=True)
assert L.number_of_nodes() == 2 * N
L.ndata['h'] = F.randn((2 * N, D))
# update node features on line graph should reflect to edge features on
# original graph.
u = [0, 0, 2, 3]
v = [1, 2, 0, 0]
eid = G.edge_ids(u, v)
L.nodes[eid].data['h'] = F.zeros((4, D))
assert F.allclose(G.edges[u, v].data['h'], F.zeros((4, D)))
# adding a new node feature on line graph should also reflect to a new
# edge feature on original graph
data = F.randn((n_edges, D))
L.ndata['w'] = data
assert F.allclose(G.edata['w'], data)
assert F.allclose(L.ndata['h'], G.edata['h'])
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_hetero_linegraph(index_dtype):
def test_line_graph2(idtype):
g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]),
'user', 'follows', index_dtype=index_dtype)
lg = dgl.line_heterograph(g)
'user', 'follows', idtype=idtype)
lg = dgl.line_graph(g)
assert lg.number_of_nodes() == 5
assert lg.number_of_edges() == 8
row, col = lg.edges()
......@@ -48,7 +38,7 @@ def test_hetero_linegraph(index_dtype):
assert np.array_equal(F.asnumpy(col),
np.array([3, 4, 0, 3, 4, 0, 1, 2]))
lg = dgl.line_heterograph(g, backtracking=False)
lg = dgl.line_graph(g, backtracking=False)
assert lg.number_of_nodes() == 5
assert lg.number_of_edges() == 4
row, col = lg.edges()
......@@ -57,8 +47,8 @@ def test_hetero_linegraph(index_dtype):
assert np.array_equal(F.asnumpy(col),
np.array([4, 0, 3, 1]))
g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]),
'user', 'follows', restrict_format='csr', index_dtype=index_dtype)
lg = dgl.line_heterograph(g)
'user', 'follows', idtype=idtype).formats('csr')
lg = dgl.line_graph(g)
assert lg.number_of_nodes() == 5
assert lg.number_of_edges() == 8
row, col = lg.edges()
......@@ -67,9 +57,9 @@ def test_hetero_linegraph(index_dtype):
assert np.array_equal(F.asnumpy(col),
np.array([3, 4, 0, 3, 4, 0, 1, 2]))
g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]),
'user', 'follows', restrict_format='csc', index_dtype=index_dtype)
lg = dgl.line_heterograph(g)
g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]),
'user', 'follows', idtype=idtype).formats('csc')
lg = dgl.line_graph(g)
assert lg.number_of_nodes() == 5
assert lg.number_of_edges() == 8
row, col, eid = lg.edges('all')
......@@ -82,6 +72,7 @@ def test_hetero_linegraph(index_dtype):
assert np.array_equal(col[order],
np.array([3, 4, 0, 3, 4, 0, 1, 2]))
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_no_backtracking():
N = 5
G = dgl.DGLGraph(nx.star_graph(N))
......@@ -94,8 +85,10 @@ def test_no_backtracking():
assert not L.has_edge_between(e2, e1)
# reverse graph related
def test_reverse():
@parametrize_dtype
def test_reverse(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(5)
# The graph need not to be completely connected.
g.add_edges([0, 1, 2], [1, 2, 1])
......@@ -113,12 +106,12 @@ def test_reverse():
assert g.edge_id(1, 2) == rg.edge_id(2, 1)
assert g.edge_id(2, 1) == rg.edge_id(1, 2)
# test dgl.reverse_heterograph
# test dgl.reverse
# test homogeneous graph
g = dgl.graph((F.tensor([0, 1, 2]), F.tensor([1, 2, 0])))
g.ndata['h'] = F.tensor([[0.], [1.], [2.]])
g.edata['h'] = F.tensor([[3.], [4.], [5.]])
g_r = dgl.reverse_heterograph(g)
g_r = dgl.reverse(g)
assert g.number_of_nodes() == g_r.number_of_nodes()
assert g.number_of_edges() == g_r.number_of_edges()
u_g, v_g, eids_g = g.all_edges(form='all')
......@@ -130,14 +123,14 @@ def test_reverse():
assert len(g_r.edata) == 0
# without share ndata
g_r = dgl.reverse_heterograph(g, copy_ndata=False)
g_r = dgl.reverse(g, copy_ndata=False)
assert g.number_of_nodes() == g_r.number_of_nodes()
assert g.number_of_edges() == g_r.number_of_edges()
assert len(g_r.ndata) == 0
assert len(g_r.edata) == 0
# with share ndata and edata
g_r = dgl.reverse_heterograph(g, copy_ndata=True, copy_edata=True)
g_r = dgl.reverse(g, copy_ndata=True, copy_edata=True)
assert g.number_of_nodes() == g_r.number_of_nodes()
assert g.number_of_edges() == g_r.number_of_edges()
assert F.array_equal(g.ndata['h'], g_r.ndata['h'])
......@@ -157,13 +150,14 @@ def test_reverse():
g = dgl.heterograph({
('user', 'follows', 'user'): ([0, 1, 2, 4, 3 ,1, 3], [1, 2, 3, 2, 0, 0, 1]),
('user', 'plays', 'game'): ([0, 0, 2, 3, 3, 4, 1], [1, 0, 1, 0, 1, 0, 0]),
('developer', 'develops', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])})
('developer', 'develops', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])},
idtype=idtype, device=F.ctx())
g.nodes['user'].data['h'] = F.tensor([0, 1, 2, 3, 4])
g.nodes['user'].data['hh'] = F.tensor([1, 1, 1, 1, 1])
g.nodes['game'].data['h'] = F.tensor([0, 1])
g.edges['follows'].data['h'] = F.tensor([0, 1, 2, 4, 3 ,1, 3])
g.edges['follows'].data['hh'] = F.tensor([1, 2, 3, 2, 0, 0, 1])
g_r = dgl.reverse_heterograph(g)
g_r = dgl.reverse(g)
for etype_g, etype_gr in zip(g.canonical_etypes, g_r.canonical_etypes):
assert etype_g[0] == etype_gr[2]
......@@ -193,7 +187,7 @@ def test_reverse():
assert F.array_equal(eids_g, eids_rg)
# withour share ndata
g_r = dgl.reverse_heterograph(g, copy_ndata=False)
g_r = dgl.reverse(g, copy_ndata=False)
for etype_g, etype_gr in zip(g.canonical_etypes, g_r.canonical_etypes):
assert etype_g[0] == etype_gr[2]
assert etype_g[1] == etype_gr[1]
......@@ -204,7 +198,7 @@ def test_reverse():
assert len(g_r.nodes['user'].data) == 0
assert len(g_r.nodes['game'].data) == 0
g_r = dgl.reverse_heterograph(g, copy_ndata=True, copy_edata=True)
g_r = dgl.reverse(g, copy_ndata=True, copy_edata=True)
print(g_r)
for etype_g, etype_gr in zip(g.canonical_etypes, g_r.canonical_etypes):
assert etype_g[0] == etype_gr[2]
......@@ -225,8 +219,10 @@ def test_reverse():
assert ('hhh' in g_r.edges['follows'].data) is True
def test_reverse_shared_frames():
@parametrize_dtype
def test_reverse_shared_frames(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(3)
g.add_edges([0, 1, 2], [1, 2, 1])
g.ndata['h'] = F.tensor([[0.], [1.], [2.]])
......@@ -238,18 +234,6 @@ def test_reverse_shared_frames():
assert F.allclose(g.edges[[0, 2], [1, 1]].data['h'],
rg.edges[[1, 1], [0, 2]].data['h'])
rg.ndata['h'] = rg.ndata['h'] + 1
assert F.allclose(rg.ndata['h'], g.ndata['h'])
g.edata['h'] = g.edata['h'] - 1
assert F.allclose(rg.edata['h'], g.edata['h'])
src_msg = fn.copy_src(src='h', out='m')
sum_reduce = fn.sum(msg='m', out='h')
rg.update_all(src_msg, sum_reduce)
assert F.allclose(g.ndata['h'], rg.ndata['h'])
def test_to_bidirected():
# homogeneous graph
g = dgl.graph((F.tensor([0, 1, 3, 1]), F.tensor([1, 2, 0, 2])))
......@@ -329,6 +313,7 @@ def test_to_bidirected():
assert F.array_equal(v, vb)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_simple_graph():
elist = [(0, 1), (0, 2), (1, 2), (0, 1)]
g = dgl.DGLGraph(elist, readonly=True)
......@@ -340,8 +325,8 @@ def test_simple_graph():
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == set(elist)
def test_bidirected_graph():
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def _test_bidirected_graph():
def _test(in_readonly, out_readonly):
elist = [(0, 0), (0, 1), (1, 0),
(1, 1), (2, 1), (2, 2)]
......@@ -361,6 +346,7 @@ def test_bidirected_graph():
_test(False, False)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_khop_graph():
N = 20
feat = F.randn((N, 5))
......@@ -386,6 +372,7 @@ def test_khop_graph():
g = dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3, directed=True))
_test(g)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_khop_adj():
N = 20
feat = F.randn((N, 5))
......@@ -402,6 +389,7 @@ def test_khop_adj():
assert F.allclose(h_0, h_1, rtol=1e-3, atol=1e-3)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_laplacian_lambda_max():
N = 20
eps = 1e-6
......@@ -410,6 +398,7 @@ def test_laplacian_lambda_max():
l_max = dgl.laplacian_lambda_max(g)
assert (l_max[0] < 2 + eps)
# test batched DGLGraph
'''
N_arr = [20, 30, 10, 12]
bg = dgl.batch([
dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3))
......@@ -419,25 +408,7 @@ def test_laplacian_lambda_max():
assert len(l_max_arr) == len(N_arr)
for l_max in l_max_arr:
assert l_max < 2 + eps
def test_add_self_loop():
g = dgl.DGLGraph()
g.add_nodes(5)
g.add_edges([0, 1, 2], [1, 1, 2])
# Nodes 0, 3, 4 don't have self-loop
new_g = dgl.transform.add_self_loop(g)
assert F.allclose(new_g.edges()[0], F.tensor([0, 0, 1, 2, 3, 4]))
assert F.allclose(new_g.edges()[1], F.tensor([1, 0, 1, 2, 3, 4]))
def test_remove_self_loop():
g = dgl.DGLGraph()
g.add_nodes(5)
g.add_edges([0, 1, 2], [1, 1, 2])
new_g = dgl.transform.remove_self_loop(g)
assert F.allclose(new_g.edges()[0], F.tensor([0]))
assert F.allclose(new_g.edges()[1], F.tensor([1]))
'''
def create_large_graph_index(num_nodes):
row = np.random.choice(num_nodes, num_nodes * 10)
......@@ -454,8 +425,9 @@ def get_nodeflow(g, node_ids, num_layers):
seed_nodes=node_ids)
return next(iter(sampler))
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_partition_with_halo():
g = dgl.DGLGraph(create_large_graph_index(1000), readonly=True)
g = dgl.DGLGraphStale(create_large_graph_index(1000), readonly=True)
node_part = np.random.choice(4, g.number_of_nodes())
subgs = dgl.transform.partition_graph_with_halo(g, node_part, 2)
for part_id, subg in subgs.items():
......@@ -484,7 +456,7 @@ def test_partition_with_halo():
@unittest.skipIf(F._default_context_str == 'gpu', reason="METIS doesn't support GPU")
def test_metis_partition():
# TODO(zhengda) Metis fails to partition a small graph.
g = dgl.DGLGraph(create_large_graph_index(1000), readonly=True)
g = dgl.DGLGraphStale(create_large_graph_index(1000), readonly=True)
check_metis_partition(g, 0)
check_metis_partition(g, 1)
check_metis_partition(g, 2)
......@@ -493,7 +465,7 @@ def test_metis_partition():
@unittest.skipIf(F._default_context_str == 'gpu', reason="METIS doesn't support GPU")
def test_hetero_metis_partition():
# TODO(zhengda) Metis fails to partition a small graph.
g = dgl.DGLGraph(create_large_graph_index(1000), readonly=True)
g = dgl.DGLGraphStale(create_large_graph_index(1000), readonly=True)
g = dgl.as_heterograph(g)
check_metis_partition(g, 0)
check_metis_partition(g, 1)
......@@ -583,7 +555,7 @@ def check_metis_partition(g, extra_hops):
@unittest.skipIf(F._default_context_str == 'gpu', reason="It doesn't support GPU")
def test_reorder_nodes():
g = dgl.DGLGraph(create_large_graph_index(1000), readonly=True)
g = dgl.DGLGraphStale(create_large_graph_index(1000), readonly=True)
new_nids = np.random.permutation(g.number_of_nodes())
# TODO(zhengda) we need to test both CSR and COO.
new_g = dgl.transform.reorder_nodes(g, new_nids)
......@@ -618,14 +590,14 @@ def test_reorder_nodes():
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_in_subgraph(index_dtype):
g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', index_dtype=index_dtype)
g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', index_dtype=index_dtype)
g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', index_dtype=index_dtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', index_dtype=index_dtype)
def test_in_subgraph(idtype):
g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', idtype=idtype)
g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', idtype=idtype)
g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', idtype=idtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', idtype=idtype)
hg = dgl.hetero_from_relations([g1, g2, g3, g4])
subg = dgl.in_subgraph(hg, {'user' : [0,1], 'game' : 0})
assert subg._idtype_str == index_dtype
assert subg.idtype == idtype
assert len(subg.ntypes) == 3
assert len(subg.etypes) == 4
u, v = subg['follow'].edges()
......@@ -644,14 +616,14 @@ def test_in_subgraph(index_dtype):
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_out_subgraph(index_dtype):
g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', index_dtype=index_dtype)
g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', index_dtype=index_dtype)
g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', index_dtype=index_dtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', index_dtype=index_dtype)
def test_out_subgraph(idtype):
g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', idtype=idtype)
g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', idtype=idtype)
g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', idtype=idtype)
g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', idtype=idtype)
hg = dgl.hetero_from_relations([g1, g2, g3, g4])
subg = dgl.out_subgraph(hg, {'user' : [0,1], 'game' : 0})
assert subg._idtype_str == index_dtype
assert subg.idtype == idtype
assert len(subg.ntypes) == 3
assert len(subg.etypes) == 4
u, v = subg['follow'].edges()
......@@ -673,20 +645,20 @@ def test_out_subgraph(index_dtype):
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented")
@parametrize_dtype
def test_compact(index_dtype):
def test_compact(idtype):
g1 = dgl.heterograph({
('user', 'follow', 'user'): [(1, 3), (3, 5)],
('user', 'plays', 'game'): [(2, 4), (3, 4), (2, 5)],
('game', 'wished-by', 'user'): [(6, 7), (5, 7)]},
{'user': 20, 'game': 10}, index_dtype=index_dtype)
{'user': 20, 'game': 10}, idtype=idtype)
g2 = dgl.heterograph({
('game', 'clicked-by', 'user'): [(3, 1)],
('user', 'likes', 'user'): [(1, 8), (8, 9)]},
{'user': 20, 'game': 10}, index_dtype=index_dtype)
{'user': 20, 'game': 10}, idtype=idtype)
g3 = dgl.graph([(0, 1), (1, 2)], num_nodes=10, ntype='user', index_dtype=index_dtype)
g4 = dgl.graph([(1, 3), (3, 5)], num_nodes=10, ntype='user', index_dtype=index_dtype)
g3 = dgl.graph([(0, 1), (1, 2)], num_nodes=10, ntype='user', idtype=idtype)
g4 = dgl.graph([(1, 3), (3, 5)], num_nodes=10, ntype='user', idtype=idtype)
def _check(g, new_g, induced_nodes):
assert g.ntypes == new_g.ntypes
......@@ -709,15 +681,15 @@ def test_compact(index_dtype):
new_g1 = dgl.compact_graphs(g1)
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g1._idtype_str == index_dtype
assert new_g1.idtype == idtype
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
assert set(induced_nodes['game']) == set([4, 5, 6])
_check(g1, new_g1, induced_nodes)
# Test with always_preserve given a dict
new_g1 = dgl.compact_graphs(
g1, always_preserve={'game': F.tensor([4, 7], dtype=getattr(F, index_dtype))})
assert new_g1._idtype_str == index_dtype
g1, always_preserve={'game': F.tensor([4, 7], idtype)})
assert new_g1.idtype == idtype
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
......@@ -726,11 +698,11 @@ def test_compact(index_dtype):
# Test with always_preserve given a tensor
new_g3 = dgl.compact_graphs(
g3, always_preserve=F.tensor([1, 7], dtype=getattr(F, index_dtype)))
g3, always_preserve=F.tensor([1, 7], idtype))
induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g3._idtype_str == index_dtype
assert new_g3.idtype == idtype
assert set(induced_nodes['user']) == set([0, 1, 2, 7])
_check(g3, new_g3, induced_nodes)
......@@ -738,8 +710,8 @@ def test_compact(index_dtype):
new_g1, new_g2 = dgl.compact_graphs([g1, g2])
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g1._idtype_str == index_dtype
assert new_g2._idtype_str == index_dtype
assert new_g1.idtype == idtype
assert new_g2.idtype == idtype
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
assert set(induced_nodes['game']) == set([3, 4, 5, 6])
_check(g1, new_g1, induced_nodes)
......@@ -747,11 +719,11 @@ def test_compact(index_dtype):
# Test multiple graphs with always_preserve given a dict
new_g1, new_g2 = dgl.compact_graphs(
[g1, g2], always_preserve={'game': F.tensor([4, 7], dtype=getattr(F, index_dtype))})
[g1, g2], always_preserve={'game': F.tensor([4, 7], dtype=idtype)})
induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in new_g1.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g1._idtype_str == index_dtype
assert new_g2._idtype_str == index_dtype
assert new_g1.idtype == idtype
assert new_g2.idtype == idtype
assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
assert set(induced_nodes['game']) == set([3, 4, 5, 6, 7])
_check(g1, new_g1, induced_nodes)
......@@ -759,19 +731,20 @@ def test_compact(index_dtype):
# Test multiple graphs with always_preserve given a tensor
new_g3, new_g4 = dgl.compact_graphs(
[g3, g4], always_preserve=F.tensor([1, 7], dtype=getattr(F, index_dtype)))
[g3, g4], always_preserve=F.tensor([1, 7], dtype=idtype))
induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in new_g3.ntypes}
induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
assert new_g3._idtype_str == index_dtype
assert new_g4._idtype_str == index_dtype
assert new_g3.idtype == idtype
assert new_g4.idtype == idtype
assert set(induced_nodes['user']) == set([0, 1, 2, 3, 5, 7])
_check(g3, new_g3, induced_nodes)
_check(g4, new_g4, induced_nodes)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU to simple not implemented")
@parametrize_dtype
def test_to_simple(index_dtype):
def test_to_simple(idtype):
# homogeneous graph
g = dgl.graph((F.tensor([0, 1, 2, 1]), F.tensor([1, 2, 0, 2])))
g.ndata['h'] = F.tensor([[0.], [1.], [2.]])
......@@ -809,7 +782,7 @@ def test_to_simple(index_dtype):
('user', 'follow', 'user'): ([0, 1, 2, 1, 1, 1],
[1, 3, 2, 3, 4, 4]),
('user', 'plays', 'game'): ([3, 2, 1, 1, 3, 2, 2], [5, 3, 4, 4, 5, 3, 3])},
index_dtype=index_dtype)
idtype=idtype, device=F.ctx())
g.nodes['user'].data['h'] = F.tensor([0, 1, 2, 3, 4])
g.nodes['user'].data['hh'] = F.tensor([0, 1, 2, 3, 4])
g.edges['follow'].data['h'] = F.tensor([0, 1, 2, 3, 4, 5])
......@@ -855,7 +828,7 @@ def test_to_simple(index_dtype):
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented")
@parametrize_dtype
def test_to_block(index_dtype):
def test_to_block(idtype):
def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True):
if dst_nodes is not None:
assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes)
......@@ -892,7 +865,7 @@ def test_to_block(index_dtype):
g = dgl.heterograph({
('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)],
('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)],
('B', 'BA', 'A'): [(2, 3), (3, 2)]}, index_dtype=index_dtype)
('B', 'BA', 'A'): [(2, 3), (3, 2)]}, idtype=idtype)
g.nodes['A'].data['x'] = F.randn((5, 10))
g.nodes['B'].data['x'] = F.randn((7, 5))
g.edges['AA'].data['x'] = F.randn((4, 3))
......@@ -929,7 +902,7 @@ def test_to_block(index_dtype):
assert bg.number_of_src_nodes() == 4
assert bg.number_of_dst_nodes() == 4
dst_nodes = F.tensor([4, 3, 2, 1], dtype=getattr(F, index_dtype))
dst_nodes = F.tensor([4, 3, 2, 1], dtype=idtype)
bg = dgl.to_block(g_a, dst_nodes)
check(g_a, bg, 'A', 'AA', dst_nodes)
check_features(g_a, bg)
......@@ -937,14 +910,14 @@ def test_to_block(index_dtype):
g_ab = g['AB']
bg = dgl.to_block(g_ab)
assert bg._idtype_str == index_dtype
assert bg.idtype == idtype
assert bg.number_of_nodes('SRC/B') == 4
assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID])
assert bg.number_of_nodes('DST/A') == 0
checkall(g_ab, bg, None)
check_features(g_ab, bg)
dst_nodes = {'B': F.tensor([5, 6, 3, 1], dtype=getattr(F, index_dtype))}
dst_nodes = {'B': F.tensor([5, 6, 3, 1], dtype=idtype)}
bg = dgl.to_block(g, dst_nodes)
assert bg.number_of_nodes('SRC/B') == 4
assert F.array_equal(bg.srcnodes['B'].data[dgl.NID], bg.dstnodes['B'].data[dgl.NID])
......@@ -952,14 +925,14 @@ def test_to_block(index_dtype):
checkall(g, bg, dst_nodes)
check_features(g, bg)
dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=getattr(F, index_dtype)), 'B': F.tensor([3, 5, 6, 1], dtype=getattr(F, index_dtype))}
dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=idtype), 'B': F.tensor([3, 5, 6, 1], dtype=idtype)}
bg = dgl.to_block(g, dst_nodes=dst_nodes)
checkall(g, bg, dst_nodes)
check_features(g, bg)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_remove_edges(index_dtype):
def test_remove_edges(idtype):
def check(g1, etype, g, edges_removed):
src, dst, eid = g.edges(etype=etype, form='all')
src1, dst1 = g1.edges(etype=etype, order='eid')
......@@ -982,82 +955,503 @@ def test_remove_edges(index_dtype):
for fmt in ['coo', 'csr', 'csc']:
for edges_to_remove in [[2], [2, 2], [3, 2], [1, 3, 1, 2]]:
g = dgl.graph([(0, 1), (2, 3), (1, 2), (3, 4)], restrict_format=fmt, index_dtype=index_dtype)
g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, getattr(F, index_dtype)))
g = dgl.graph([(0, 1), (2, 3), (1, 2), (3, 4)], idtype=idtype).formats(fmt)
g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, idtype))
check(g1, None, g, edges_to_remove)
g = dgl.graph(
spsp.csr_matrix(([1, 1, 1, 1], ([0, 2, 1, 3], [1, 3, 2, 4])), shape=(5, 5)),
restrict_format=fmt, index_dtype=index_dtype)
g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, getattr(F, index_dtype)))
idtype=idtype).formats(fmt)
g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, idtype))
check(g1, None, g, edges_to_remove)
g = dgl.heterograph({
('A', 'AA', 'A'): [(0, 1), (2, 3), (1, 2), (3, 4)],
('A', 'AB', 'B'): [(0, 1), (1, 3), (3, 5), (1, 6)],
('B', 'BA', 'A'): [(2, 3), (3, 2)]}, index_dtype=index_dtype)
g2 = dgl.remove_edges(g, {'AA': F.tensor([2], getattr(F, index_dtype)), 'AB': F.tensor([3], getattr(F, index_dtype)), 'BA': F.tensor([1], getattr(F, index_dtype))})
('B', 'BA', 'A'): [(2, 3), (3, 2)]}, idtype=idtype)
g2 = dgl.remove_edges(g, {'AA': F.tensor([2], idtype), 'AB': F.tensor([3], idtype), 'BA': F.tensor([1], idtype)})
check(g2, 'AA', g, [2])
check(g2, 'AB', g, [3])
check(g2, 'BA', g, [1])
g3 = dgl.remove_edges(g, {'AA': F.tensor([], getattr(F, index_dtype)), 'AB': F.tensor([3], getattr(F, index_dtype)), 'BA': F.tensor([1], getattr(F, index_dtype))})
g3 = dgl.remove_edges(g, {'AA': F.tensor([], idtype), 'AB': F.tensor([3], idtype), 'BA': F.tensor([1], idtype)})
check(g3, 'AA', g, [])
check(g3, 'AB', g, [3])
check(g3, 'BA', g, [1])
g4 = dgl.remove_edges(g, {'AB': F.tensor([3, 1, 2, 0], getattr(F, index_dtype))})
g4 = dgl.remove_edges(g, {'AB': F.tensor([3, 1, 2, 0], idtype)})
check(g4, 'AA', g, [])
check(g4, 'AB', g, [3, 1, 2, 0])
check(g4, 'BA', g, [])
def test_cast():
m = spsp.coo_matrix(([1, 1], ([0, 1], [1, 2])), (4, 4))
g = dgl.DGLGraph(m, readonly=True)
gsrc, gdst = g.edges(order='eid')
ndata = F.randn((4, 5))
edata = F.randn((2, 4))
g.ndata['x'] = ndata
g.edata['y'] = edata
hg = dgl.as_heterograph(g, 'A', 'AA')
assert hg.ntypes == ['A']
assert hg.etypes == ['AA']
assert hg.canonical_etypes == [('A', 'AA', 'A')]
assert hg.number_of_nodes() == 4
assert hg.number_of_edges() == 2
hgsrc, hgdst = hg.edges(order='eid')
assert F.array_equal(gsrc, hgsrc)
assert F.array_equal(gdst, hgdst)
g2 = dgl.as_immutable_graph(hg)
assert g2.number_of_nodes() == 4
assert g2.number_of_edges() == 2
g2src, g2dst = hg.edges(order='eid')
assert F.array_equal(g2src, gsrc)
assert F.array_equal(g2dst, gdst)
@parametrize_dtype
def test_add_edges(idtype):
# homogeneous graph
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
u = 0
v = 1
g = dgl.add_edges(g, u, v)
assert g.device == F.ctx()
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 3
u = [0]
v = [1]
g = dgl.add_edges(g, u, v)
assert g.device == F.ctx()
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 4
u = F.tensor(u, dtype=idtype)
v = F.tensor(v, dtype=idtype)
g = dgl.add_edges(g, u, v)
assert g.device == F.ctx()
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 5
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))
# node id larger than current max node id
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
u = F.tensor([0, 1], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g = dgl.add_edges(g, u, v)
assert g.number_of_nodes() == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
# has data
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
u = F.tensor([0, 1], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g = dgl.add_edges(g, u, v, e_feat)
assert g.number_of_nodes() == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))
# zero data graph
g = dgl.graph([], num_nodes=0, idtype=idtype, device=F.ctx())
u = F.tensor([0, 1], dtype=idtype)
v = F.tensor([2, 2], dtype=idtype)
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g = dgl.add_edges(g, u, v, e_feat)
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 2
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([2, 2], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([2, 2], dtype=idtype))
assert F.array_equal(g.edata['hh'], F.tensor([2, 2], dtype=idtype))
# bipartite graph
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
u = 0
v = 1
g = dgl.add_edges(g, u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 3
u = [0]
v = [1]
g = dgl.add_edges(g, u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 4
u = F.tensor(u, dtype=idtype)
v = F.tensor(v, dtype=idtype)
g = dgl.add_edges(g, u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 5
u, v = g.edges(form='uv')
assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))
# node id larger than current max node id
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g = dgl.add_edges(g, u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
# has data
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
g.ndata['h'] = {'user' : F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx()),
'game' : F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())}
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g = dgl.add_edges(g, u, v, e_feat)
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 0], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 0], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))
# heterogeneous graph
g = create_test_heterograph4(idtype)
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g = dgl.add_edges(g, u, v, etype='plays')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_nodes('developer') == 2
assert g.number_of_edges('plays') == 6
assert g.number_of_edges('develops') == 2
u, v = g.edges(form='uv', order='eid', etype='plays')
assert F.array_equal(u, F.tensor([0, 1, 1, 2, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 0, 1, 1, 2, 3], dtype=idtype))
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 0, 0], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 1, 1, 1, 0, 0], dtype=idtype))
# add with feature
e_feat = {'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 1, 1], dtype=idtype), ctx=F.ctx())
g = dgl.add_edges(g, u, v, data=e_feat, etype='develops')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_nodes('developer') == 3
assert g.number_of_edges('plays') == 6
assert g.number_of_edges('develops') == 4
u, v = g.edges(form='uv', order='eid', etype='develops')
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 2, 3], dtype=idtype))
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3, 0], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 1, 1], dtype=idtype))
assert F.array_equal(g.edges['develops'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype))
@parametrize_dtype
def test_add_nodes(idtype):
# homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1,1,1], dtype=idtype), ctx=F.ctx())
new_g = dgl.add_nodes(g, 1)
assert g.number_of_nodes() == 3
assert new_g.number_of_nodes() == 4
assert F.array_equal(new_g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
# zero node graph
g = dgl.graph([], num_nodes=3, idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1,1,1], dtype=idtype), ctx=F.ctx())
g = dgl.add_nodes(g, 1, data={'h' : F.copy_to(F.tensor([2], dtype=idtype), ctx=F.ctx())})
assert g.number_of_nodes() == 4
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 2], dtype=idtype))
# bipartite graph
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
g = dgl.add_nodes(g, 2, data={'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, ntype='user')
assert g.number_of_nodes('user') == 4
assert g.number_of_nodes('game') == 3
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype))
g = dgl.add_nodes(g, 2, ntype='game')
assert g.number_of_nodes('user') == 4
assert g.number_of_nodes('game') == 5
# heterogeneous graph
g = create_test_heterograph4(idtype)
g = dgl.add_nodes(g, 1, ntype='user')
g = dgl.add_nodes(g, 2, data={'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, ntype='game')
assert g.number_of_nodes('user') == 4
assert g.number_of_nodes('game') == 4
assert g.number_of_nodes('developer') == 2
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 2], dtype=idtype))
@parametrize_dtype
def test_remove_edges(idtype):
# homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
e = 0
g = dgl.remove_edges(g, e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
e = [0]
g = dgl.remove_edges(g, e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
e = F.tensor([0], dtype=idtype)
g = dgl.remove_edges(g, e)
assert g.number_of_edges() == 0
# has node data
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g = dgl.remove_edges(g, 1)
assert g.number_of_edges() == 1
assert F.array_equal(g.ndata['h'], F.tensor([1, 2, 3], dtype=idtype))
# has edge data
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
g = dgl.remove_edges(g, 0)
assert g.number_of_edges() == 1
assert F.array_equal(g.edata['h'], F.tensor([2], dtype=idtype))
# invalid eid
assert_fail = False
try:
g = dgl.remove_edges(g, 1)
except:
assert_fail = True
assert assert_fail
# bipartite graph
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
e = 0
g = dgl.remove_edges(g, e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
e = [0]
g = dgl.remove_edges(g, e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
e = F.tensor([0], dtype=idtype)
g = dgl.remove_edges(g, e)
assert g.number_of_edges() == 0
# has data
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
g.ndata['h'] = {'user' : F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx()),
'game' : F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())}
g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
g = dgl.remove_edges(g, 1)
assert g.number_of_edges() == 1
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([1], dtype=idtype))
# heterogeneous graph
g = create_test_heterograph4(idtype)
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
g = dgl.remove_edges(g, 1, etype='plays')
assert g.number_of_edges('plays') == 3
u, v = g.edges(form='uv', order='eid', etype='plays')
assert F.array_equal(u, F.tensor([0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 1], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 3, 4], dtype=idtype))
# remove all edges of 'develops'
g = dgl.remove_edges(g, [0, 1], etype='develops')
assert g.number_of_edges('develops') == 0
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2], dtype=idtype))
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3], dtype=idtype))
@parametrize_dtype
def test_remove_nodes(idtype):
# homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
n = 0
g = dgl.remove_nodes(g, n)
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
n = [1]
g = dgl.remove_nodes(g, n)
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 0
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
n = F.tensor([2], dtype=idtype)
g = dgl.remove_nodes(g, n)
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
# invalid nid
assert_fail = False
try:
g.remove_nodes(3)
except:
assert_fail = True
assert assert_fail
# has node and edge data
g = dgl.graph(([0, 0, 2], [0, 1, 2]), idtype=idtype, device=F.ctx())
g.ndata['hv'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g = dgl.remove_nodes(g, F.tensor([0], dtype=idtype))
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
assert F.array_equal(g.ndata['hv'], F.tensor([2, 3], dtype=idtype))
assert F.array_equal(g.edata['he'], F.tensor([3], dtype=idtype))
# node id larger than current max node id
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
n = 0
g = dgl.remove_nodes(g, n, ntype='user')
assert g.number_of_nodes('user') == 1
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
n = [1]
g = dgl.remove_nodes(g, n, ntype='user')
assert g.number_of_nodes('user') == 1
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
n = F.tensor([0], dtype=idtype)
g = dgl.remove_nodes(g, n, ntype='game')
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 2
assert g.number_of_edges() == 2
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([0 ,1], dtype=idtype))
# heterogeneous graph
g = create_test_heterograph4(idtype)
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
g = dgl.remove_nodes(g, 0, ntype='game')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 1
assert g.number_of_nodes('developer') == 2
assert g.number_of_edges('plays') == 2
assert g.number_of_edges('develops') == 1
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2], dtype=idtype))
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3], dtype=idtype))
u, v = g.edges(form='uv', order='eid', etype='plays')
assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 0], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([3, 4], dtype=idtype))
u, v = g.edges(form='uv', order='eid', etype='develops')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([0], dtype=idtype))
@parametrize_dtype
def test_add_selfloop(idtype):
# homogeneous graph
g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g = dgl.add_self_loop(g)
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 6
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edata['he'], F.tensor([1, 2, 3, 0, 0, 0], dtype=idtype))
# bipartite graph
g = dgl.bipartite(([0, 1, 2], [1, 2, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
# nothing will happend
raise_error = False
try:
g = dgl.add_self_loop(g)
except:
raise_error = True
assert raise_error
g = create_test_heterograph6(idtype)
g = dgl.add_self_loop(g, etype='follows')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 2
assert g.number_of_edges('follows') == 5
assert g.number_of_edges('plays') == 2
u, v = g.edges(form='uv', order='eid', etype='follows')
assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1, 2, 0, 0, 0], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 2], dtype=idtype))
raise_error = False
try:
g = dgl.add_self_loop(g, etype='plays')
except:
raise_error = True
assert raise_error
@parametrize_dtype
def test_remove_selfloop(idtype):
# homogeneous graph
g = dgl.graph(([0, 0, 0, 1], [1, 0, 0, 2]), idtype=idtype, device=F.ctx())
g.edata['he'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
g = dgl.remove_self_loop(g)
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 2
assert F.array_equal(g.edata['he'], F.tensor([1, 4], dtype=idtype))
# bipartite graph
g = dgl.bipartite(([0, 1, 2], [1, 2, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
# nothing will happend
raise_error = False
try:
g = dgl.remove_self_loop(g, etype='plays')
except:
raise_error = True
assert raise_error
g = create_test_heterograph5(idtype)
g = dgl.remove_self_loop(g, etype='follows')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 2
assert g.number_of_edges('follows') == 2
assert g.number_of_edges('plays') == 2
u, v = g.edges(form='uv', order='eid', etype='follows')
assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h'], F.tensor([2, 4], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 2], dtype=idtype))
raise_error = False
try:
g = dgl.remove_self_loop(g, etype='plays')
except:
raise_error = True
assert raise_error
if __name__ == '__main__':
# test_reorder_nodes()
# test_line_graph()
# test_no_backtracking()
# test_reverse()
# test_reverse_shared_frames()
# test_to_bidirected()
# test_simple_graph()
# test_bidirected_graph()
# test_khop_adj()
# test_khop_graph()
# test_laplacian_lambda_max()
# test_remove_self_loop()
# test_add_self_loop()
# test_partition_with_halo()
test_metis_partition()
test_hetero_metis_partition()
# test_hetero_linegraph('int32')
# test_compact()
# test_to_simple("int32")
# test_in_subgraph("int32")
# test_out_subgraph()
# test_to_block("int32")
# test_remove_edges()
pass
tests/compute/test_traversal.py
View file @ 44089c8b
import random
import sys
import time
import unittest
import dgl
import networkx as nx
......@@ -17,8 +18,9 @@ def toset(x):
# F.zerocopy_to_numpy may return a int
return set(F.zerocopy_to_numpy(x).tolist())
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_bfs(index_dtype, n=100):
def test_bfs(idtype, n=100):
def _bfs_nx(g_nx, src):
edges = nx.bfs_edges(g_nx, src)
layers_nx = [set([src])]
......@@ -28,25 +30,20 @@ def test_bfs(index_dtype, n=100):
for u, v in edges:
if u in layers_nx[-1]:
frontier.add(v)
edge_frontier.add(g.edge_id(u, v))
edge_frontier.add(g.edge_ids(u, v))
else:
layers_nx.append(frontier)
edges_nx.append(edge_frontier)
frontier = set([v])
edge_frontier = set([g.edge_id(u, v)])
edge_frontier = set([g.edge_ids(u, v)])
# avoids empty successors
if len(frontier) > 0 and len(edge_frontier) > 0:
layers_nx.append(frontier)
edges_nx.append(edge_frontier)
return layers_nx, edges_nx
g = dgl.DGLGraph()
a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
g.from_scipy_sparse_matrix(a)
if index_dtype == 'int32':
g = dgl.graph(g.edges()).int()
else:
g = dgl.graph(g.edges()).long()
g = dgl.graph(a).astype(idtype)
g_nx = g.to_networkx()
src = random.choice(range(n))
......@@ -56,29 +53,19 @@ def test_bfs(index_dtype, n=100):
assert all(toset(x) == y for x, y in zip(layers_dgl, layers_nx))
g_nx = nx.random_tree(n, seed=42)
g = dgl.DGLGraph()
g.from_networkx(g_nx)
if index_dtype == 'int32':
g = dgl.graph(g.edges()).int()
else:
g = dgl.graph(g.edges()).long()
g = dgl.graph(g_nx).astype(idtype)
src = 0
_, edges_nx = _bfs_nx(g_nx, src)
edges_dgl = dgl.bfs_edges_generator(g, src)
assert len(edges_dgl) == len(edges_nx)
assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_topological_nodes(index_dtype, n=100):
g = dgl.DGLGraph()
def test_topological_nodes(idtype, n=100):
a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
b = sp.tril(a, -1).tocoo()
g.from_scipy_sparse_matrix(b)
if index_dtype == 'int32':
g = dgl.graph(g.edges()).int()
else:
g = dgl.graph(g.edges()).long()
g = dgl.graph(b).astype(idtype)
layers_dgl = dgl.topological_nodes_generator(g)
......@@ -101,15 +88,12 @@ def test_topological_nodes(index_dtype, n=100):
assert all(toset(x) == toset(y) for x, y in zip(layers_dgl, layers_spmv))
DFS_LABEL_NAMES = ['forward', 'reverse', 'nontree']
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
@parametrize_dtype
def test_dfs_labeled_edges(index_dtype, example=False):
dgl_g = dgl.DGLGraph()
def test_dfs_labeled_edges(idtype, example=False):
dgl_g = dgl.DGLGraph().astype(idtype)
dgl_g.add_nodes(6)
dgl_g.add_edges([0, 1, 0, 3, 3], [1, 2, 2, 4, 5])
if index_dtype == 'int32':
dgl_g = dgl.graph(dgl_g.edges()).int()
else:
dgl_g = dgl.graph(dgl_g.edges()).long()
dgl_edges, dgl_labels = dgl.dfs_labeled_edges_generator(
dgl_g, [0, 3], has_reverse_edge=True, has_nontree_edge=True)
dgl_edges = [toset(t) for t in dgl_edges]
......@@ -141,6 +125,6 @@ def test_dfs_labeled_edges(index_dtype, example=False):
assert False
if __name__ == '__main__':
test_bfs(index_dtype='int32')
test_topological_nodes(index_dtype='int32')
test_dfs_labeled_edges(index_dtype='int32')
test_bfs(idtype='int32')
test_topological_nodes(idtype='int32')
test_dfs_labeled_edges(idtype='int32')
tests/compute/test_udf.py deleted 100644 → 0
View file @ 015acfd2
import backend as F
import dgl
import networkx as nx
import dgl.utils as utils
from dgl import DGLGraph, ALL
from dgl.udf import NodeBatch, EdgeBatch
def test_node_batch():
g = dgl.DGLGraph(nx.path_graph(20))
feat = F.randn((g.number_of_nodes(), 10))
g.ndata['x'] = feat
# test all
v = utils.toindex(slice(0, g.number_of_nodes()))
n_repr = g.get_n_repr(v)
nbatch = NodeBatch(v, n_repr)
assert F.allclose(nbatch.data['x'], feat)
assert nbatch.mailbox is None
assert F.allclose(nbatch.nodes(), g.nodes())
assert nbatch.batch_size() == g.number_of_nodes()
assert len(nbatch) == g.number_of_nodes()
# test partial
v = utils.toindex(F.tensor([0, 3, 5, 7, 9]))
n_repr = g.get_n_repr(v)
nbatch = NodeBatch(v, n_repr)
assert F.allclose(nbatch.data['x'], F.gather_row(feat, F.tensor([0, 3, 5, 7, 9])))
assert nbatch.mailbox is None
assert F.allclose(nbatch.nodes(), F.tensor([0, 3, 5, 7, 9]))
assert nbatch.batch_size() == 5
assert len(nbatch) == 5
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\
F.shape(ebatch.dst['x'])[1] == d
assert F.shape(ebatch.data['x'])[0] == g.number_of_edges() and\
F.shape(ebatch.data['x'])[1] == d
assert F.allclose(ebatch.edges()[0], u.tousertensor())
assert F.allclose(ebatch.edges()[1], v.tousertensor())
assert F.allclose(ebatch.edges()[2], F.arange(0, g.number_of_edges()))
assert ebatch.batch_size() == g.number_of_edges()
assert len(ebatch) == g.number_of_edges()
# test partial
eid = utils.toindex(F.tensor([0, 3, 5, 7, 11, 13, 15, 27]))
u, v, _ = g._graph.find_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] == 8 and\
F.shape(ebatch.src['x'])[1] == d
assert F.shape(ebatch.dst['x'])[0] == 8 and\
F.shape(ebatch.dst['x'])[1] == d
assert F.shape(ebatch.data['x'])[0] == 8 and\
F.shape(ebatch.data['x'])[1] == d
assert F.allclose(ebatch.edges()[0], u.tousertensor())
assert F.allclose(ebatch.edges()[1], v.tousertensor())
assert F.allclose(ebatch.edges()[2], eid.tousertensor())
assert ebatch.batch_size() == 8
assert len(ebatch) == 8
if __name__ == '__main__':
test_node_batch()
test_edge_batch()
tests/compute/utils.py
View file @ 44089c8b
import pytest
parametrize_dtype = pytest.mark.parametrize("index_dtype", ['int32', 'int64'])
import backend as F
if F._default_context_str == 'cpu':
parametrize_dtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
else:
# only test int32 on GPU because many graph operators are not supported for int64.
parametrize_dtype = pytest.mark.parametrize("idtype", [F.int32])
def check_fail(fn, *args, **kwargs):
try:
......
tests/cpp/test_aten.cc
View file @ 44089c8b
......@@ -189,19 +189,23 @@ TEST(ArrayTest, TestArith) {
};
template <typename IDX>
void _TestHStack() {
IdArray a = aten::Range(0, 100, sizeof(IDX)*8, CTX);
IdArray b = aten::Range(100, 200, sizeof(IDX)*8, CTX);
IdArray c = aten::HStack(a, b);
void _TestHStack(DLContext ctx) {
IdArray a = aten::Range(0, 100, sizeof(IDX)*8, ctx);
IdArray b = aten::Range(100, 200, sizeof(IDX)*8, ctx);
IdArray c = aten::HStack(a, b).CopyTo(aten::CPU);
ASSERT_EQ(c->ndim, 1);
ASSERT_EQ(c->shape[0], 200);
for (int i = 0; i < 200; ++i)
ASSERT_EQ(Ptr<IDX>(c)[i], i);
}
TEST(ArrayTest, TestHStack) {
_TestHStack<int32_t>();
_TestHStack<int64_t>();
TEST(ArrayTest, HStack) {
_TestHStack<int32_t>(CPU);
_TestHStack<int64_t>(CPU);
#ifdef DGL_USE_CUDA
_TestHStack<int32_t>(GPU);
_TestHStack<int64_t>(GPU);
#endif
}
template <typename IDX>
......@@ -1238,6 +1242,61 @@ TEST(ArrayTest, CumSum) {
#endif
}
template <typename IDX, typename D>
void _TestScatter_(DLContext ctx) {
IdArray out = aten::Full(1, 10, 8*sizeof(IDX), ctx);
IdArray idx = aten::VecToIdArray(std::vector<IDX>({2, 3, 9}), sizeof(IDX)*8, ctx);
IdArray val = aten::VecToIdArray(std::vector<IDX>({-20, 30, 90}), sizeof(IDX)*8, ctx);
aten::Scatter_(idx, val, out);
IdArray tout = aten::VecToIdArray(std::vector<IDX>({1, 1, -20, 30, 1, 1, 1, 1, 1, 90}), sizeof(IDX)*8, ctx);
ASSERT_TRUE(ArrayEQ<IDX>(out, tout));
}
TEST(ArrayTest, Scatter_) {
_TestScatter_<int32_t, int32_t>(CPU);
_TestScatter_<int64_t, int32_t>(CPU);
_TestScatter_<int32_t, int64_t>(CPU);
_TestScatter_<int64_t, int64_t>(CPU);
#ifdef DGL_USE_CUDA
_TestScatter_<int32_t, int32_t>(GPU);
_TestScatter_<int64_t, int32_t>(GPU);
_TestScatter_<int32_t, int64_t>(GPU);
_TestScatter_<int64_t, int64_t>(GPU);
#endif
}
template <typename IDX>
void _TestNonZero(DLContext ctx) {
auto val = aten::VecToIdArray(std::vector<IDX>({0, 1, 2, 0, -10, 0, 0, 23}), sizeof(IDX)*8, ctx);
auto idx = aten::NonZero(val);
auto tidx = aten::VecToIdArray(std::vector<int64_t>({1, 2, 4, 7}), 64, ctx);
ASSERT_TRUE(ArrayEQ<IDX>(idx, tidx));
val = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX)*8, ctx);
idx = aten::NonZero(val);
tidx = aten::VecToIdArray(std::vector<int64_t>({}), 64, ctx);
ASSERT_TRUE(ArrayEQ<IDX>(idx, tidx));
val = aten::VecToIdArray(std::vector<IDX>({0, 0, 0, 0}), sizeof(IDX)*8, ctx);
idx = aten::NonZero(val);
tidx = aten::VecToIdArray(std::vector<int64_t>({}), 64, ctx);
ASSERT_TRUE(ArrayEQ<IDX>(idx, tidx));
val = aten::Full(1, 3, sizeof(IDX)*8, ctx);
idx = aten::NonZero(val);
tidx = aten::VecToIdArray(std::vector<int64_t>({0, 1, 2}), 64, ctx);
ASSERT_TRUE(ArrayEQ<IDX>(idx, tidx));
}
TEST(ArrayTest, NonZero) {
_TestNonZero<int32_t>(CPU);
_TestNonZero<int64_t>(CPU);
#ifdef DGL_USE_CUDA
_TestNonZero<int32_t>(GPU);
_TestNonZero<int64_t>(GPU);
#endif
}
template <typename IdType>
void _TestLineGraphCOO(DLContext ctx) {
/*
......@@ -1344,16 +1403,3 @@ TEST(LineGraphTest, LineGraphCOO) {
_TestLineGraphCOO<int32_t>(CPU);
_TestLineGraphCOO<int64_t>(CPU);
}
template <typename IDX>
void _TestNonZero() {
BoolArray a = aten::VecToIdArray(std::vector<IDX>({1, 0, 1, 1, 0, 0, 1}));
IdArray indices = aten::NonZero(a);
IdArray expected = aten::VecToIdArray(std::vector<IDX>({0, 2, 3, 6}));
ASSERT_TRUE(ArrayEQ<IDX>(indices, expected));
}
TEST(ArrayTest, NonZero) {
_TestNonZero<int32_t>();
_TestNonZero<int64_t>();
}
tests/cpp/test_serialize.cc
View file @ 44089c8b
......@@ -4,6 +4,7 @@
#include <gtest/gtest.h>
#include <algorithm>
#include <iostream>
#include <memory>
#include <vector>
#include "../../src/graph/heterograph.h"
#include "../../src/graph/unit_graph.h"
......@@ -18,7 +19,7 @@ TEST(Serialize, UnitGraph_COO) {
auto src = VecToIdArray<int64_t>({1, 2, 5, 3});
auto dst = VecToIdArray<int64_t>({1, 6, 2, 6});
auto mg = std::dynamic_pointer_cast<UnitGraph>(
dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst, dgl::SparseFormat::kCOO));
dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst, coo_code));
std::string blob;
dmlc::MemoryStringStream ifs(&blob);
......@@ -39,13 +40,15 @@ TEST(Serialize, UnitGraph_CSR) {
aten::CSRMatrix csr_matrix;
auto src = VecToIdArray<int64_t>({1, 2, 5, 3});
auto dst = VecToIdArray<int64_t>({1, 6, 2, 6});
auto mg = std::dynamic_pointer_cast<UnitGraph>(
dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst, dgl::SparseFormat::kCSR));
auto coo_g = std::dynamic_pointer_cast<UnitGraph>(
dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst));
auto csr_g =
std::dynamic_pointer_cast<UnitGraph>(coo_g->GetGraphInFormat(csr_code));
std::string blob;
dmlc::MemoryStringStream ifs(&blob);
static_cast<dmlc::Stream *>(&ifs)->Write(mg);
static_cast<dmlc::Stream *>(&ifs)->Write(csr_g);
dmlc::MemoryStringStream ofs(&blob);
auto ug2 = Serializer::make_shared<UnitGraph>();
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment