Merge pull request #4 from jermainewang/master

Sync with latest commit

src/runtime/workspace_pool.h

@@ -3,8 +3,8 @@
  * \file workspace_pool.h
  * \brief Workspace pool utility.
  */
-#ifndef TVM_RUNTIME_WORKSPACE_POOL_H_
-#define TVM_RUNTIME_WORKSPACE_POOL_H_
+#ifndef DGL_RUNTIME_WORKSPACE_POOL_H_
+#define DGL_RUNTIME_WORKSPACE_POOL_H_
 
 #include <dgl/runtime/device_api.h>
 #include <vector>
@@ -58,4 +58,4 @@ class WorkspacePool {
 
 }  // namespace runtime
 }  // namespace tvm
-#endif  // TVM_RUNTIME_WORKSPACE_POOL_H_
+#endif  // DGL_RUNTIME_WORKSPACE_POOL_H_

src/scheduler/scheduler.cc

-// DGL Scheduler implementation
-
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file scheduler/scheduler.cc
+ * \brief DGL Scheduler implementation
+ */
+#include <dgl/scheduler.h>
 #include <unordered_map>
 #include <vector>
-#include <dgl/scheduler.h>
 
 namespace dgl {
 namespace sched {
@@ -19,7 +22,7 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
 
     // bkt: deg->dsts
     std::unordered_map<int64_t, std::vector<int64_t>> bkt;
-    for (auto& it: in_edges) {
+    for (const auto& it : in_edges) {
         bkt[it.second.size()].push_back(it.first);
     }
 
@@ -38,15 +41,15 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
     int64_t* msec_ptr = static_cast<int64_t*>(mid_section->data);
 
     // fill in bucketing ordering
-    for (auto& it: bkt) { // for each bucket
-        int64_t deg = it.first;
-        int64_t n_dst = it.second.size();
+    for (const auto& it : bkt) {  // for each bucket
+        const int64_t deg = it.first;
+        const int64_t n_dst = it.second.size();
         *deg_ptr++ = deg;
         *nsec_ptr++ = n_dst;
         *msec_ptr++ = deg * n_dst;
-        for (auto dst: it.second) { // for each dst in this bucket
+        for (const auto dst : it.second) {  // for each dst in this bucket
             *nid_ptr++ = dst;
-            for (auto mid: in_edges[dst]) { // for each in edge of dst
+            for (const auto mid : in_edges[dst]) {  // for each in edge of dst
                 *mid_ptr++ = mid;
             }
         }
@@ -62,6 +65,6 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
     return std::move(ret);
 }
 
-} // namespace sched
+}  // namespace sched
 
-} // namespace dgl
+}  // namespace dgl

src/scheduler/scheduler_apis.cc

-#include "../c_api_common.h"
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file scheduler/scheduler_apis.cc
+ * \brief DGL scheduler APIs
+ */
 #include <dgl/graph.h>
 #include <dgl/scheduler.h>
+#include "../c_api_common.h"
 
 using tvm::runtime::TVMArgs;
 using tvm::runtime::TVMRetValue;
@@ -18,8 +23,8 @@ TVM_REGISTER_GLOBAL("scheduler._CAPI_DGLDegreeBucketingFromGraph")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    auto edges = gptr->Edges(false);
+    const auto& edges = gptr->Edges(false);
     *rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(edges.dst));
   });
 
-} // namespace dgl
+}  // namespace dgl

tests/mxnet/test_basics.py

@@ -9,7 +9,7 @@ reduce_msg_shapes = set()
 
 def check_eq(a, b):
     assert a.shape == b.shape
-    assert mx.sum(a == b) == int(np.prod(list(a.shape)))
+    assert mx.nd.sum(a == b).asnumpy() == int(np.prod(list(a.shape)))
 
 def message_func(src, edge):
     assert len(src['h'].shape) == 2
@@ -53,16 +53,12 @@ def test_batch_setter_getter():
     assert len(g.get_n_repr()) == 0
     g.set_n_repr({'h' : mx.nd.zeros((10, D))})
     # set partial nodes
-    # TODO we need to enable the test later.
-    '''
     u = mx.nd.array([1, 3, 5], dtype='int64')
     g.set_n_repr({'h' : mx.nd.ones((3, D))}, u)
     assert _pfc(g.get_n_repr()['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
     # get partial nodes
     u = mx.nd.array([1, 2, 3], dtype='int64')
-    print(g.get_n_repr(u)['h'])
     assert _pfc(g.get_n_repr(u)['h']) == [1., 0., 1.]
-    '''
 
     '''
     s, d, eid
@@ -127,9 +123,11 @@ def test_batch_setter_autograd():
     with mx.autograd.record():
         g = generate_graph(grad=True)
         h1 = g.get_n_repr()['h']
+        h1.attach_grad()
         # partial set
         v = mx.nd.array([1, 2, 8], dtype='int64')
         hh = mx.nd.zeros((len(v), D))
+        hh.attach_grad()
         g.set_n_repr({'h' : hh}, v)
         h2 = g.get_n_repr()['h']
     h2.backward(mx.nd.ones((10, D)) * 2)
@@ -252,8 +250,7 @@ def test_pull_0deg():
 
 if __name__ == '__main__':
     test_batch_setter_getter()
-    # TODO we need to enable it after index_copy is implemented.
-    #test_batch_setter_autograd()
+    test_batch_setter_autograd()
     test_batch_send()
     test_batch_recv()
     test_update_routines()

tests/pytorch/test_basics.py

@@ -20,22 +20,26 @@ def reduce_func(node, msgs):
     reduce_msg_shapes.add(tuple(msgs.shape))
     assert len(msgs.shape) == 3
     assert msgs.shape[2] == D
-    return {'m' : th.sum(msgs, 1)}
+    return {'accum' : th.sum(msgs, 1)}
 
 def apply_node_func(node):
-    return {'h' : node['h'] + node['m']}
+    return {'h' : node['h'] + node['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
+    # 17 edges
     for i in range(1, 9):
         g.add_edge(0, i)
         g.add_edge(i, 9)
     # add a back flow from 9 to 0
     g.add_edge(9, 0)
     ncol = Variable(th.randn(10, D), requires_grad=grad)
+    accumcol = Variable(th.randn(10, D), requires_grad=grad)
+    ecol = Variable(th.randn(17, D), requires_grad=grad)
     g.set_n_repr({'h' : ncol})
+    g.set_n_initializer(lambda shape, dtype : th.zeros(shape))
     return g
 
 def test_batch_setter_getter():
@@ -46,8 +50,9 @@ def test_batch_setter_getter():
     g.set_n_repr({'h' : th.zeros((10, D))})
     assert _pfc(g.get_n_repr()['h']) == [0.] * 10
     # pop nodes
+    old_len = len(g.get_n_repr())
     assert _pfc(g.pop_n_repr('h')) == [0.] * 10
-    assert len(g.get_n_repr()) == 0
+    assert len(g.get_n_repr()) == old_len - 1
     g.set_n_repr({'h' : th.zeros((10, D))})
     # set partial nodes
     u = th.tensor([1, 3, 5])
@@ -81,8 +86,9 @@ def test_batch_setter_getter():
     g.set_e_repr({'l' : th.zeros((17, D))})
     assert _pfc(g.get_e_repr()['l']) == [0.] * 17
     # pop edges
+    old_len = len(g.get_e_repr())
     assert _pfc(g.pop_e_repr('l')) == [0.] * 17
-    assert len(g.get_e_repr()) == 0
+    assert len(g.get_e_repr()) == old_len - 1
     g.set_e_repr({'l' : th.zeros((17, D))})
     # set partial edges (many-many)
     u = th.tensor([0, 0, 2, 5, 9])
@@ -203,14 +209,13 @@ def test_reduce_0deg():
     g.add_edge(3, 0)
     g.add_edge(4, 0)
     def _message(src, edge):
-        return src
+        return {'m' : src['h']}
     def _reduce(node, msgs):
-        assert msgs is not None
-        return node + msgs.sum(1)
+        return {'h' : node['h'] + msgs['m'].sum(1)}
     old_repr = th.randn(5, 5)
-    g.set_n_repr(old_repr)
+    g.set_n_repr({'h' : old_repr})
     g.update_all(_message, _reduce)
-    new_repr = g.get_n_repr()
+    new_repr = g.get_n_repr()['h']
 
     assert th.allclose(new_repr[1:], old_repr[1:])
     assert th.allclose(new_repr[0], old_repr.sum(0))
@@ -220,29 +225,30 @@ def test_pull_0deg():
     g.add_nodes(2)
     g.add_edge(0, 1)
     def _message(src, edge):
-        return src
+        return {'m' : src['h']}
     def _reduce(node, msgs):
-        assert msgs is not None
-        return msgs.sum(1)
-
+        return {'h' : msgs['m'].sum(1)}
     old_repr = th.randn(2, 5)
-    g.set_n_repr(old_repr)
+    g.set_n_repr({'h' : old_repr})
+
     g.pull(0, _message, _reduce)
-    new_repr = g.get_n_repr()
+    new_repr = g.get_n_repr()['h']
     assert th.allclose(new_repr[0], old_repr[0])
     assert th.allclose(new_repr[1], old_repr[1])
+
     g.pull(1, _message, _reduce)
-    new_repr = g.get_n_repr()
+    new_repr = g.get_n_repr()['h']
     assert th.allclose(new_repr[1], old_repr[0])
 
     old_repr = th.randn(2, 5)
-    g.set_n_repr(old_repr)
+    g.set_n_repr({'h' : old_repr})
     g.pull([0, 1], _message, _reduce)
-    new_repr = g.get_n_repr()
+    new_repr = g.get_n_repr()['h']
     assert th.allclose(new_repr[0], old_repr[0])
     assert th.allclose(new_repr[1], old_repr[0])
 
-def test_send_twice():
+def _disabled_test_send_twice():
+    # TODO(minjie): please re-enable this unittest after the send code problem is fixed.
     g = DGLGraph()
     g.add_nodes(3)
     g.add_edge(0, 1)
@@ -348,5 +354,4 @@ if __name__ == '__main__':
     test_update_routines()
     test_reduce_0deg()
     test_pull_0deg()
-    test_send_twice()
     test_send_multigraph()

tests/pytorch/test_basics_anonymous.py

-import torch as th
-from torch.autograd import Variable
-import numpy as np
-from dgl.graph import DGLGraph, __REPR__
-
-D = 32
-reduce_msg_shapes = set()
-
-def check_eq(a, b):
-    assert a.shape == b.shape
-    assert th.sum(a == b) == int(np.prod(list(a.shape)))
-
-def message_func(hu, e_uv):
-    assert len(hu.shape) == 2
-    assert hu.shape[1] == D
-    return hu
-
-def reduce_func(hv, msgs):
-    reduce_msg_shapes.add(tuple(msgs.shape))
-    assert len(msgs.shape) == 3
-    assert msgs.shape[2] == D
-    return hv + th.sum(msgs, 1)
-
-def generate_graph(grad=False):
-    g = DGLGraph()
-    g.add_nodes(10)
-    # create a graph where 0 is the source and 9 is the sink
-    for i in range(1, 9):
-        g.add_edge(0, i)
-        g.add_edge(i, 9)
-    # add a back flow from 9 to 0
-    g.add_edge(9, 0)
-    col = Variable(th.randn(10, D), requires_grad=grad)
-    g.set_n_repr(col)
-    return g
-
-def test_batch_setter_getter():
-    def _pfc(x):
-        return list(x.numpy()[:,0])
-    g = generate_graph()
-    # set all nodes
-    g.set_n_repr(th.zeros((10, D)))
-    assert _pfc(g.get_n_repr()) == [0.] * 10
-    # pop nodes
-    assert _pfc(g.pop_n_repr()) == [0.] * 10
-    assert len(g.get_n_repr()) == 0
-    g.set_n_repr(th.zeros((10, D)))
-    # set partial nodes
-    u = th.tensor([1, 3, 5])
-    g.set_n_repr(th.ones((3, D)), u)
-    assert _pfc(g.get_n_repr()) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
-    # get partial nodes
-    u = th.tensor([1, 2, 3])
-    assert _pfc(g.get_n_repr(u)) == [1., 0., 1.]
-
-    '''
-    s, d, eid
-    0, 1, 0
-    1, 9, 1
-    0, 2, 2
-    2, 9, 3
-    0, 3, 4
-    3, 9, 5
-    0, 4, 6
-    4, 9, 7
-    0, 5, 8
-    5, 9, 9
-    0, 6, 10
-    6, 9, 11
-    0, 7, 12
-    7, 9, 13
-    0, 8, 14
-    8, 9, 15
-    9, 0, 16
-    '''
-    # set all edges
-    g.set_e_repr(th.zeros((17, D)))
-    assert _pfc(g.get_e_repr()) == [0.] * 17
-    # pop edges
-    assert _pfc(g.pop_e_repr()) == [0.] * 17
-    assert len(g.get_e_repr()) == 0
-    g.set_e_repr(th.zeros((17, D)))
-    # set partial edges (many-many)
-    u = th.tensor([0, 0, 2, 5, 9])
-    v = th.tensor([1, 3, 9, 9, 0])
-    g.set_e_repr(th.ones((5, D)), u, v)
-    truth = [0.] * 17
-    truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
-    assert _pfc(g.get_e_repr()) == truth
-    # set partial edges (many-one)
-    u = th.tensor([3, 4, 6])
-    v = th.tensor([9])
-    g.set_e_repr(th.ones((3, D)), u, v)
-    truth[5] = truth[7] = truth[11] = 1.
-    assert _pfc(g.get_e_repr()) == truth
-    # set partial edges (one-many)
-    u = th.tensor([0])
-    v = th.tensor([4, 5, 6])
-    g.set_e_repr(th.ones((3, D)), u, v)
-    truth[6] = truth[8] = truth[10] = 1.
-    assert _pfc(g.get_e_repr()) == truth
-    # get partial edges (many-many)
-    u = th.tensor([0, 6, 0])
-    v = th.tensor([6, 9, 7])
-    assert _pfc(g.get_e_repr(u, v)) == [1., 1., 0.]
-    # get partial edges (many-one)
-    u = th.tensor([5, 6, 7])
-    v = th.tensor([9])
-    assert _pfc(g.get_e_repr(u, v)) == [1., 1., 0.]
-    # get partial edges (one-many)
-    u = th.tensor([0])
-    v = th.tensor([3, 4, 5])
-    assert _pfc(g.get_e_repr(u, v)) == [1., 1., 1.]
-
-def test_batch_setter_autograd():
-    g = generate_graph(grad=True)
-    h1 = g.get_n_repr()
-    # partial set
-    v = th.tensor([1, 2, 8])
-    hh = Variable(th.zeros((len(v), D)), requires_grad=True)
-    g.set_n_repr(hh, v)
-    h2 = g.get_n_repr()
-    h2.backward(th.ones((10, D)) * 2)
-    check_eq(h1.grad[:,0], th.tensor([2., 0., 0., 2., 2., 2., 2., 2., 0., 2.]))
-    check_eq(hh.grad[:,0], th.tensor([2., 2., 2.]))
-
-def test_batch_send():
-    g = generate_graph()
-    def _fmsg(hu, edge):
-        assert hu.shape == (5, D)
-        return hu
-    g.register_message_func(_fmsg)
-    # many-many send
-    u = th.tensor([0, 0, 0, 0, 0])
-    v = th.tensor([1, 2, 3, 4, 5])
-    g.send(u, v)
-    # one-many send
-    u = th.tensor([0])
-    v = th.tensor([1, 2, 3, 4, 5])
-    g.send(u, v)
-    # many-one send
-    u = th.tensor([1, 2, 3, 4, 5])
-    v = th.tensor([9])
-    g.send(u, v)
-
-def test_batch_recv():
-    g = generate_graph()
-    g.register_message_func(message_func)
-    g.register_reduce_func(reduce_func)
-    u = th.tensor([0, 0, 0, 4, 5, 6])
-    v = th.tensor([1, 2, 3, 9, 9, 9])
-    reduce_msg_shapes.clear()
-    g.send(u, v)
-    g.recv(th.unique(v))
-    assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
-    reduce_msg_shapes.clear()
-
-def test_update_routines():
-    g = generate_graph()
-    g.register_message_func(message_func)
-    g.register_reduce_func(reduce_func)
-
-    # send_and_recv
-    reduce_msg_shapes.clear()
-    u = th.tensor([0, 0, 0, 4, 5, 6])
-    v = th.tensor([1, 2, 3, 9, 9, 9])
-    g.send_and_recv(u, v)
-    assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
-    reduce_msg_shapes.clear()
-
-    # pull
-    v = th.tensor([1, 2, 3, 9])
-    reduce_msg_shapes.clear()
-    g.pull(v)
-    assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
-    reduce_msg_shapes.clear()
-
-    # push
-    v = th.tensor([0, 1, 2, 3])
-    reduce_msg_shapes.clear()
-    g.push(v)
-    assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
-    reduce_msg_shapes.clear()
-
-    # update_all
-    reduce_msg_shapes.clear()
-    g.update_all()
-    assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
-    reduce_msg_shapes.clear()
-
-if __name__ == '__main__':
-    test_batch_setter_getter()
-    test_batch_setter_autograd()
-    test_batch_send()
-    test_batch_recv()
-    test_update_routines()

tests/pytorch/test_batched_graph.py

@@ -18,8 +18,8 @@ def tree1():
     g.add_edge(4, 1)
     g.add_edge(1, 0)
     g.add_edge(2, 0)
-    g.set_n_repr(th.Tensor([0, 1, 2, 3, 4]))
-    g.set_e_repr(th.randn(4, 10))
+    g.set_n_repr({'h' : th.Tensor([0, 1, 2, 3, 4])})
+    g.set_e_repr({'h' : th.randn(4, 10)})
     return g
 
 def tree2():
@@ -37,17 +37,17 @@ def tree2():
     g.add_edge(0, 4)
     g.add_edge(4, 1)
     g.add_edge(3, 1)
-    g.set_n_repr(th.Tensor([0, 1, 2, 3, 4]))
-    g.set_e_repr(th.randn(4, 10))
+    g.set_n_repr({'h' : th.Tensor([0, 1, 2, 3, 4])})
+    g.set_e_repr({'h' : th.randn(4, 10)})
     return g
 
 def test_batch_unbatch():
     t1 = tree1()
     t2 = tree2()
-    n1 = t1.get_n_repr()
-    n2 = t2.get_n_repr()
-    e1 = t1.get_e_repr()
-    e2 = t2.get_e_repr()
+    n1 = t1.get_n_repr()['h']
+    n2 = t2.get_n_repr()['h']
+    e1 = t1.get_e_repr()['h']
+    e2 = t2.get_e_repr()['h']
 
     bg = dgl.batch([t1, t2])
     assert bg.number_of_nodes() == 10
@@ -57,10 +57,10 @@ def test_batch_unbatch():
     assert bg.batch_num_edges == [4, 4]
 
     tt1, tt2 = dgl.unbatch(bg)
-    assert th.allclose(t1.get_n_repr(), tt1.get_n_repr())
-    assert th.allclose(t1.get_e_repr(), tt1.get_e_repr())
-    assert th.allclose(t2.get_n_repr(), tt2.get_n_repr())
-    assert th.allclose(t2.get_e_repr(), tt2.get_e_repr())
+    assert th.allclose(t1.get_n_repr()['h'], tt1.get_n_repr()['h'])
+    assert th.allclose(t1.get_e_repr()['h'], tt1.get_e_repr()['h'])
+    assert th.allclose(t2.get_n_repr()['h'], tt2.get_n_repr()['h'])
+    assert th.allclose(t2.get_e_repr()['h'], tt2.get_e_repr()['h'])
 
 def test_batch_unbatch1():
     t1 = tree1()
@@ -74,20 +74,20 @@ def test_batch_unbatch1():
     assert b2.batch_num_edges == [4, 4, 4]
 
     s1, s2, s3 = dgl.unbatch(b2)
-    assert th.allclose(t2.get_n_repr(), s1.get_n_repr())
-    assert th.allclose(t2.get_e_repr(), s1.get_e_repr())
-    assert th.allclose(t1.get_n_repr(), s2.get_n_repr())
-    assert th.allclose(t1.get_e_repr(), s2.get_e_repr())
-    assert th.allclose(t2.get_n_repr(), s3.get_n_repr())
-    assert th.allclose(t2.get_e_repr(), s3.get_e_repr())
+    assert th.allclose(t2.get_n_repr()['h'], s1.get_n_repr()['h'])
+    assert th.allclose(t2.get_e_repr()['h'], s1.get_e_repr()['h'])
+    assert th.allclose(t1.get_n_repr()['h'], s2.get_n_repr()['h'])
+    assert th.allclose(t1.get_e_repr()['h'], s2.get_e_repr()['h'])
+    assert th.allclose(t2.get_n_repr()['h'], s3.get_n_repr()['h'])
+    assert th.allclose(t2.get_e_repr()['h'], s3.get_e_repr()['h'])
 
 def test_batch_sendrecv():
     t1 = tree1()
     t2 = tree2()
 
     bg = dgl.batch([t1, t2])
-    bg.register_message_func(lambda src, edge: src)
-    bg.register_reduce_func(lambda node, msgs: th.sum(msgs, 1))
+    bg.register_message_func(lambda src, edge: {'m' : src['h']})
+    bg.register_reduce_func(lambda node, msgs: {'h' : th.sum(msgs['m'], 1)})
     u = [3, 4, 2 + 5, 0 + 5]
     v = [1, 1, 4 + 5, 4 + 5]
 
@@ -95,8 +95,8 @@ def test_batch_sendrecv():
     bg.recv(v)
 
     t1, t2 = dgl.unbatch(bg)
-    assert t1.get_n_repr()[1] == 7
-    assert t2.get_n_repr()[4] == 2
+    assert t1.get_n_repr()['h'][1] == 7
+    assert t2.get_n_repr()['h'][4] == 2
 
 
 def test_batch_propagate():
@@ -104,8 +104,8 @@ def test_batch_propagate():
     t2 = tree2()
 
     bg = dgl.batch([t1, t2])
-    bg.register_message_func(lambda src, edge: src)
-    bg.register_reduce_func(lambda node, msgs: th.sum(msgs, 1))
+    bg.register_message_func(lambda src, edge: {'m' : src['h']})
+    bg.register_reduce_func(lambda node, msgs: {'h' : th.sum(msgs['m'], 1)})
     # get leaves.
 
     order = []
@@ -123,45 +123,37 @@ def test_batch_propagate():
     bg.propagate(traverser=order)
     t1, t2 = dgl.unbatch(bg)
 
-    assert t1.get_n_repr()[0] == 9
-    assert t2.get_n_repr()[1] == 5
+    assert t1.get_n_repr()['h'][0] == 9
+    assert t2.get_n_repr()['h'][1] == 5
 
 def test_batched_edge_ordering():
     g1 = dgl.DGLGraph()
     g1.add_nodes(6)
     g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
     e1 = th.randn(5, 10)
-    g1.set_e_repr(e1)
+    g1.set_e_repr({'h' : e1})
     g2 = dgl.DGLGraph()
     g2.add_nodes(6)
     g2.add_edges([0, 1 ,2 ,5, 4 ,5], [1, 2, 3, 4, 3, 0])
     e2 = th.randn(6, 10)
-    g2.set_e_repr(e2)
+    g2.set_e_repr({'h' : e2})
     g = dgl.batch([g1, g2])
-    r1 = g.get_e_repr()[g.edge_id(4, 5)]
-    r2 = g1.get_e_repr()[g1.edge_id(4, 5)]
+    r1 = g.get_e_repr()['h'][g.edge_id(4, 5)]
+    r2 = g1.get_e_repr()['h'][g1.edge_id(4, 5)]
     assert th.equal(r1, r2)
 
 def test_batch_no_edge():
-    # FIXME: current impl cannot handle this case!!!
-    #        comment out for now to test CI
-    return
-    """
     g1 = dgl.DGLGraph()
     g1.add_nodes(6)
     g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
     e1 = th.randn(5, 10)
-    g1.set_e_repr(e1)
     g2 = dgl.DGLGraph()
     g2.add_nodes(6)
     g2.add_edges([0, 1, 2, 5, 4, 5], [1 ,2 ,3, 4, 3, 0])
     e2 = th.randn(6, 10)
-    g2.set_e_repr(e2)
     g3 = dgl.DGLGraph()
     g3.add_nodes(1)  # no edges
-
     g = dgl.batch([g1, g3, g2]) # should not throw an error
-    """
 
 if __name__ == '__main__':
     test_batch_unbatch()

tests/pytorch/test_frame.py

@@ -2,14 +2,11 @@ import torch as th
 from torch.autograd import Variable
 import numpy as np
 from dgl.frame import Frame, FrameRef
-from dgl.utils import Index
+from dgl.utils import Index, toindex
 
 N = 10
 D = 5
 
-def check_eq(a, b):
-    return a.shape == b.shape and np.allclose(a.numpy(), b.numpy())
-
 def check_fail(fn):
     try:
         fn()
@@ -27,12 +24,13 @@ def test_create():
     data = create_test_data()
     f1 = Frame()
     for k, v in data.items():
-        f1.add_column(k, v)
-    assert f1.schemes == set(data.keys())
+        f1.update_column(k, v)
+    print(f1.schemes)
+    assert f1.keys() == set(data.keys())
     assert f1.num_columns == 3
     assert f1.num_rows == N
     f2 = Frame(data)
-    assert f2.schemes == set(data.keys())
+    assert f2.keys() == set(data.keys())
     assert f2.num_columns == 3
     assert f2.num_rows == N
     f1.clear()
@@ -45,9 +43,9 @@ def test_column1():
     f = Frame(data)
     assert f.num_rows == N
     assert len(f) == 3
-    assert check_eq(f['a1'], data['a1'])
+    assert th.allclose(f['a1'].data, data['a1'].data)
     f['a1'] = data['a2']
-    assert check_eq(f['a2'], data['a2'])
+    assert th.allclose(f['a2'].data, data['a2'].data)
     # add a different length column should fail
     def failed_add_col():
         f['a4'] = th.zeros([N+1, D])
@@ -70,16 +68,15 @@ def test_column2():
     f = FrameRef(data, [3, 4, 5, 6, 7])
     assert f.num_rows == 5
     assert len(f) == 3
-    assert check_eq(f['a1'], data['a1'][3:8])
+    assert th.allclose(f['a1'], data['a1'].data[3:8])
     # set column should reflect on the referenced data
     f['a1'] = th.zeros([5, D])
-    assert check_eq(data['a1'][3:8], th.zeros([5, D]))
-    # add new column should be padded with zero
-    f['a4'] = th.ones([5, D])
-    assert len(data) == 4
-    assert check_eq(data['a4'][0:3], th.zeros([3, D]))
-    assert check_eq(data['a4'][3:8], th.ones([5, D]))
-    assert check_eq(data['a4'][8:10], th.zeros([2, D]))
+    assert th.allclose(data['a1'].data[3:8], th.zeros([5, D]))
+    # add new partial column should fail with error initializer
+    f.set_initializer(lambda shape, dtype : assert_(False))
+    def failed_add_col():
+        f['a4'] = th.ones([5, D])
+    assert check_fail(failed_add_col)
 
 def test_append1():
     # test append API on Frame
@@ -91,9 +88,14 @@ def test_append1():
     f1.append(f2)
     assert f1.num_rows == 2 * N
     c1 = f1['a1']
-    assert c1.shape == (2 * N, D)
+    assert c1.data.shape == (2 * N, D)
     truth = th.cat([data['a1'], data['a1']])
-    assert check_eq(truth, c1)
+    assert th.allclose(truth, c1.data)
+    # append dict of different length columns should fail
+    f3 = {'a1' : th.zeros((3, D)), 'a2' : th.zeros((3, D)), 'a3' : th.zeros((2, D))}
+    def failed_append():
+        f1.append(f3)
+    assert check_fail(failed_append)
 
 def test_append2():
     # test append on FrameRef
@@ -113,7 +115,7 @@ def test_append2():
     assert not f.is_span_whole_column()
     assert f.num_rows == 3 * N
     new_idx = list(range(N)) + list(range(2*N, 4*N))
-    assert check_eq(f.index().tousertensor(), th.tensor(new_idx))
+    assert th.all(f.index().tousertensor() == th.tensor(new_idx, dtype=th.int64))
     assert data.num_rows == 4 * N
 
 def test_row1():
@@ -127,13 +129,13 @@ def test_row1():
     rows = f[rowid]
     for k, v in rows.items():
         assert v.shape == (len(rowid), D)
-        assert check_eq(v, data[k][rowid])
+        assert th.allclose(v, data[k][rowid])
     # test duplicate keys
     rowid = Index(th.tensor([8, 2, 2, 1]))
     rows = f[rowid]
     for k, v in rows.items():
         assert v.shape == (len(rowid), D)
-        assert check_eq(v, data[k][rowid])
+        assert th.allclose(v, data[k][rowid])
 
     # setter
     rowid = Index(th.tensor([0, 2, 4]))
@@ -143,12 +145,14 @@ def test_row1():
             }
     f[rowid] = vals
     for k, v in f[rowid].items():
-        assert check_eq(v, th.zeros((len(rowid), D)))
+        assert th.allclose(v, th.zeros((len(rowid), D)))
 
-    # setting rows with new column should automatically add a new column
-    vals['a4'] = th.ones((len(rowid), D))
-    f[rowid] = vals
-    assert len(f) == 4
+    # setting rows with new column should raise error with error initializer
+    f.set_initializer(lambda shape, dtype : assert_(False))
+    def failed_update_rows():
+        vals['a4'] = th.ones((len(rowid), D))
+        f[rowid] = vals
+    assert check_fail(failed_update_rows)
 
 def test_row2():
     # test row getter/setter autograd compatibility
@@ -161,13 +165,13 @@ def test_row2():
     rowid = Index(th.tensor([0, 2]))
     rows = f[rowid]
     rows['a1'].backward(th.ones((len(rowid), D)))
-    assert check_eq(c1.grad[:,0], th.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
+    assert th.allclose(c1.grad[:,0], th.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
     c1.grad.data.zero_()
     # test duplicate keys
     rowid = Index(th.tensor([8, 2, 2, 1]))
     rows = f[rowid]
     rows['a1'].backward(th.ones((len(rowid), D)))
-    assert check_eq(c1.grad[:,0], th.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))
+    assert th.allclose(c1.grad[:,0], th.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))
     c1.grad.data.zero_()
 
     # setter
@@ -180,8 +184,8 @@ def test_row2():
     f[rowid] = vals
     c11 = f['a1']
     c11.backward(th.ones((N, D)))
-    assert check_eq(c1.grad[:,0], th.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
-    assert check_eq(vals['a1'].grad, th.ones((len(rowid), D)))
+    assert th.allclose(c1.grad[:,0], th.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
+    assert th.allclose(vals['a1'].grad, th.ones((len(rowid), D)))
     assert vals['a2'].grad is None
 
 def test_row3():
@@ -201,8 +205,9 @@ def test_row3():
     newidx = list(range(N))
     newidx.pop(2)
     newidx.pop(2)
+    newidx = toindex(newidx)
     for k, v in f.items():
-        assert check_eq(v, data[k][th.tensor(newidx)])
+        assert th.allclose(v, data[k][newidx])
 
 def test_sharing():
     data = Frame(create_test_data())
@@ -210,10 +215,10 @@ def test_sharing():
     f2 = FrameRef(data, index=[2, 3, 4, 5, 6])
     # test read
     for k, v in f1.items():
-        assert check_eq(data[k][0:4], v)
+        assert th.allclose(data[k].data[0:4], v)
     for k, v in f2.items():
-        assert check_eq(data[k][2:7], v)
-    f2_a1 = f2['a1']
+        assert th.allclose(data[k].data[2:7], v)
+    f2_a1 = f2['a1'].data
     # test write
     # update own ref should not been seen by the other.
     f1[Index(th.tensor([0, 1]))] = {
@@ -221,7 +226,7 @@ def test_sharing():
             'a2' : th.zeros([2, D]),
             'a3' : th.zeros([2, D]),
             }
-    assert check_eq(f2['a1'], f2_a1)
+    assert th.allclose(f2['a1'], f2_a1)
     # update shared space should been seen by the other.
     f1[Index(th.tensor([2, 3]))] = {
             'a1' : th.ones([2, D]),
@@ -229,7 +234,7 @@ def test_sharing():
             'a3' : th.ones([2, D]),
             }
     f2_a1[0:2] = th.ones([2, D])
-    assert check_eq(f2['a1'], f2_a1)
+    assert th.allclose(f2['a1'], f2_a1)
 
 if __name__ == '__main__':
     test_create()

tests/pytorch/test_function.py

 import torch as th
 import dgl
 import dgl.function as fn
-from dgl.graph import __REPR__
 
 def generate_graph():
     g = dgl.DGLGraph()
@@ -37,18 +36,9 @@ def generate_graph1():
     g.set_e_repr(h)
     return g
 
-def reducer_msg(node, msgs):
-    return th.sum(msgs['m'], 1)
-
-def reducer_out(node, msgs):
-    return {'h' : th.sum(msgs, 1)}
-
 def reducer_both(node, msgs):
     return {'h' : th.sum(msgs['m'], 1)}
 
-def reducer_none(node, msgs):
-    return th.sum(msgs, 1)
-
 def test_copy_src():
     # copy_src with both fields
     g = generate_graph()
@@ -58,30 +48,6 @@ def test_copy_src():
     assert th.allclose(g.get_n_repr()['h'],
             th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
 
-    # copy_src with only src field; the out field should use anonymous repr
-    g = generate_graph()
-    g.register_message_func(fn.copy_src(src='h'))
-    g.register_reduce_func(reducer_out)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
-
-    # copy_src with no src field; should use anonymous repr
-    g = generate_graph1()
-    g.register_message_func(fn.copy_src(out='m'))
-    g.register_reduce_func(reducer_both)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
-
-    # copy src with no fields;
-    g = generate_graph1()
-    g.register_message_func(fn.copy_src())
-    g.register_reduce_func(reducer_out)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
-
 def test_copy_edge():
     # copy_edge with both fields
     g = generate_graph()
@@ -91,30 +57,6 @@ def test_copy_edge():
     assert th.allclose(g.get_n_repr()['h'],
             th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
 
-    # copy_edge with only edge field; the out field should use anonymous repr
-    g = generate_graph()
-    g.register_message_func(fn.copy_edge(edge='h'))
-    g.register_reduce_func(reducer_out)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
-
-    # copy_edge with no edge field; should use anonymous repr
-    g = generate_graph1()
-    g.register_message_func(fn.copy_edge(out='m'))
-    g.register_reduce_func(reducer_both)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
-
-    # copy edge with no fields;
-    g = generate_graph1()
-    g.register_message_func(fn.copy_edge())
-    g.register_reduce_func(reducer_out)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
-
 def test_src_mul_edge():
     # src_mul_edge with all fields
     g = generate_graph()
@@ -124,34 +66,6 @@ def test_src_mul_edge():
     assert th.allclose(g.get_n_repr()['h'],
             th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
 
-    g = generate_graph()
-    g.register_message_func(fn.src_mul_edge(src='h', edge='h'))
-    g.register_reduce_func(reducer_out)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
-
-    g = generate_graph1()
-    g.register_message_func(fn.src_mul_edge(out='m'))
-    g.register_reduce_func(reducer_both)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
-
-    g = generate_graph1()
-    g.register_message_func(fn.src_mul_edge())
-    g.register_reduce_func(reducer_out)
-    g.update_all()
-    assert th.allclose(g.get_n_repr()['h'],
-            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
-
-    g = generate_graph1()
-    g.register_message_func(fn.src_mul_edge())
-    g.register_reduce_func(reducer_none)
-    g.update_all()
-    assert th.allclose(g.get_n_repr(),
-            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
-
 if __name__ == '__main__':
     test_copy_src()
     test_copy_edge()

tests/pytorch/test_line_graph.py

@@ -5,35 +5,31 @@ import dgl
 
 D = 5
 
-def check_eq(a, b):
-    return a.shape == b.shape and np.allclose(a.numpy(), b.numpy())
-
 def test_line_graph():
     N = 5
     G = dgl.DGLGraph(nx.star_graph(N))
-    G.set_e_repr(th.randn((2 * N, D)))
+    G.set_e_repr({'h' : th.randn((2 * N, D))})
     n_edges = G.number_of_edges()
     L = G.line_graph(shared=True)
     assert L.number_of_nodes() == 2 * N
-    L.set_n_repr(th.randn((2 * N, D)))
+    L.set_n_repr({'h' : th.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.set_n_repr(th.zeros((4, D)), eid)
-    assert check_eq(G.get_e_repr(u, v), th.zeros((4, D)))
+    L.set_n_repr({'h' : th.zeros((4, D))}, eid)
+    assert th.allclose(G.get_e_repr(u, v)['h'], th.zeros((4, D)))
 
     # adding a new node feature on line graph should also reflect to a new
     # edge feature on original graph
     data = th.randn(n_edges, D)
     L.set_n_repr({'w': data})
-    assert check_eq(G.get_e_repr()['w'], data)
+    assert th.allclose(G.get_e_repr()['w'], data)
 
 def test_no_backtracking():
     N = 5
     G = dgl.DGLGraph(nx.star_graph(N))
-    G.set_e_repr(th.randn((2 * N, D)))
     L = G.line_graph(backtracking=False)
     assert L.number_of_nodes() == 2 * N
     for i in range(1, N):

tests/pytorch/test_specialization.py

@@ -22,23 +22,23 @@ def generate_graph():
 def test_update_all():
     def _test(fld):
         def message_func(hu, edge):
-            return hu[fld]
+            return {'m' : hu[fld]}
 
         def message_func_edge(hu, edge):
             if len(hu[fld].shape) == 1:
-                return hu[fld] * edge['e1']
+                return {'m' : hu[fld] * edge['e1']}
             else:
-                return hu[fld] * edge['e2']
+                return {'m' : hu[fld] * edge['e2']}
 
         def reduce_func(hv, msgs):
-            return {fld : th.sum(msgs, 1)}
+            return {fld : th.sum(msgs['m'], 1)}
 
         def apply_func(hu):
             return {fld : 2 * hu[fld]}
         g = generate_graph()
         # update all
         v1 = g.get_n_repr()[fld]
-        g.update_all(fn.copy_src(src=fld), fn.sum(out=fld), apply_func)
+        g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
         v2 = g.get_n_repr()[fld]
         g.set_n_repr({fld : v1})
         g.update_all(message_func, reduce_func, apply_func)
@@ -46,12 +46,12 @@ def test_update_all():
         assert th.allclose(v2, v3)
         # update all with edge weights
         v1 = g.get_n_repr()[fld]
-        g.update_all(fn.src_mul_edge(src=fld, edge='e1'),
-                fn.sum(out=fld), apply_func)
+        g.update_all(fn.src_mul_edge(src=fld, edge='e1', out='m'),
+                fn.sum(msg='m', out=fld), apply_func)
         v2 = g.get_n_repr()[fld]
         g.set_n_repr({fld : v1})
-        g.update_all(fn.src_mul_edge(src=fld, edge='e2'),
-                fn.sum(out=fld), apply_func)
+        g.update_all(fn.src_mul_edge(src=fld, edge='e2', out='m'),
+                fn.sum(msg='m', out=fld), apply_func)
         v3 = g.get_n_repr()[fld]
         g.set_n_repr({fld : v1})
         g.update_all(message_func_edge, reduce_func, apply_func)
@@ -68,42 +68,40 @@ def test_send_and_recv():
     v = th.tensor([1, 2, 3, 9, 9, 0])
     def _test(fld):
         def message_func(hu, edge):
-            return hu[fld]
+            return {'m' : hu[fld]}
 
         def message_func_edge(hu, edge):
             if len(hu[fld].shape) == 1:
-                return hu[fld] * edge['e1']
+                return {'m' : hu[fld] * edge['e1']}
             else:
-                return hu[fld] * edge['e2']
+                return {'m' : hu[fld] * edge['e2']}
 
         def reduce_func(hv, msgs):
-            return {fld : th.sum(msgs, 1)}
+            return {fld : th.sum(msgs['m'], 1)}
 
         def apply_func(hu):
             return {fld : 2 * hu[fld]}
         g = generate_graph()
         # send and recv
         v1 = g.get_n_repr()[fld]
-        g.send_and_recv(u, v, fn.copy_src(src=fld),
-                fn.sum(out=fld), apply_func)
+        g.send_and_recv(u, v, fn.copy_src(src=fld, out='m'),
+                fn.sum(msg='m', out=fld), apply_func)
         v2 = g.get_n_repr()[fld]
         g.set_n_repr({fld : v1})
-        g.send_and_recv(u, v, message_func,
-                reduce_func, apply_func)
+        g.send_and_recv(u, v, message_func, reduce_func, apply_func)
         v3 = g.get_n_repr()[fld]
         assert th.allclose(v2, v3)
         # send and recv with edge weights
         v1 = g.get_n_repr()[fld]
-        g.send_and_recv(u, v, fn.src_mul_edge(src=fld, edge='e1'),
-                fn.sum(out=fld), apply_func)
+        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.get_n_repr()[fld]
         g.set_n_repr({fld : v1})
-        g.send_and_recv(u, v, fn.src_mul_edge(src=fld, edge='e2'),
-                fn.sum(out=fld), apply_func)
+        g.send_and_recv(u, v, fn.src_mul_edge(src=fld, edge='e2', out='m'),
+                fn.sum(msg='m', out=fld), apply_func)
         v3 = g.get_n_repr()[fld]
         g.set_n_repr({fld : v1})
-        g.send_and_recv(u, v, message_func_edge,
-                reduce_func, apply_func)
+        g.send_and_recv(u, v, message_func_edge, reduce_func, apply_func)
         v4 = g.get_n_repr()[fld]
         assert th.allclose(v2, v3)
         assert th.allclose(v3, v4)
@@ -123,22 +121,23 @@ def test_update_all_multi_fn():
         return {'v2': th.sum(msgs['m2'], 1)}
 
     g = generate_graph()
+    g.set_n_repr({'v1' : th.zeros((10,)), 'v2' : th.zeros((10,))})
     fld = 'f2'
     # update all, mix of builtin and UDF
     g.update_all([fn.copy_src(src=fld, out='m1'), message_func],
-                 [fn.sum(msgs='m1', out='v1'), reduce_func],
+                 [fn.sum(msg='m1', out='v1'), reduce_func],
                  None)
     v1 = g.get_n_repr()['v1']
     v2 = g.get_n_repr()['v2']
     assert th.allclose(v1, v2)
 
     # run builtin with single message and reduce
-    g.update_all(fn.copy_src(src=fld), fn.sum(out='v1'), None)
+    g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'), None)
     v1 = g.get_n_repr()['v1']
     assert th.allclose(v1, v2)
 
-    # 1 message, 2 reduces, using anonymous repr
-    g.update_all(fn.copy_src(src=fld), [fn.sum(out='v2'), fn.sum(out='v3')], None)
+    # 1 message, 2 reduces
+    g.update_all(fn.copy_src(src=fld, out='m'), [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')], None)
     v2 = g.get_n_repr()['v2']
     v3 = g.get_n_repr()['v3']
     assert th.allclose(v1, v2)
@@ -146,7 +145,7 @@ def test_update_all_multi_fn():
 
     # update all with edge weights, 2 message, 3 reduces
     g.update_all([fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
-                 [fn.sum(msgs='m1', out='v1'), fn.sum(msgs='m2', out='v2'), fn.sum(msgs='m1', out='v3')],
+                 [fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
                  None)
     v1 = g.get_n_repr()['v1']
     v2 = g.get_n_repr()['v2']
@@ -173,25 +172,30 @@ def test_send_and_recv_multi_fn():
         return {'v2' : th.sum(msgs['m2'], 1)}
 
     g = generate_graph()
+    g.set_n_repr({'v1' : th.zeros((10, D)), 'v2' : th.zeros((10, D)),
+        'v3' : th.zeros((10, D))})
     fld = 'f2'
 
     # send and recv, mix of builtin and UDF
     g.send_and_recv(u, v,
                     [fn.copy_src(src=fld, out='m1'), message_func],
-                    [fn.sum(msgs='m1', out='v1'), reduce_func],
+                    [fn.sum(msg='m1', out='v1'), reduce_func],
                     None)
     v1 = g.get_n_repr()['v1']
     v2 = g.get_n_repr()['v2']
     assert th.allclose(v1, v2)
 
     # run builtin with single message and reduce
-    g.send_and_recv(u, v, fn.copy_src(src=fld), fn.sum(out='v1'),
+    g.send_and_recv(u, v, fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'),
                     None)
     v1 = g.get_n_repr()['v1']
     assert th.allclose(v1, v2)
 
-    # 1 message, 2 reduces, using anonymous repr
-    g.send_and_recv(u, v, fn.copy_src(src=fld), [fn.sum(out='v2'), fn.sum(out='v3')], None)
+    # 1 message, 2 reduces
+    g.send_and_recv(u, v,
+            fn.copy_src(src=fld, out='m'),
+            [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')],
+            None)
     v2 = g.get_n_repr()['v2']
     v3 = g.get_n_repr()['v3']
     assert th.allclose(v1, v2)
@@ -200,7 +204,7 @@ def test_send_and_recv_multi_fn():
     # send and recv with edge weights, 2 message, 3 reduces
     g.send_and_recv(u, v,
                     [fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
-                    [fn.sum(msgs='m1', out='v1'), fn.sum(msgs='m2', out='v2'), fn.sum(msgs='m1', out='v3')],
+                    [fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
                     None)
     v1 = g.get_n_repr()['v1']
     v2 = g.get_n_repr()['v2']

tests/scripts/task_lint.sh

+#!/bin/sh
+
+# cpplint
+echo 'Checking code style of C++ codes...'
+python3 third_party/dmlc-core/scripts/lint.py dgl cpp include src

tests/scripts/test_examples.sh

+#!/bin/bash
+
+GCN_EXAMPLE_DIR="../../examples/pytorch/gcn"
+
+function fail {
+    echo FAIL: $@
+    exit -1
+}
+
+function usage {
+    echo "Usage: $0 [CPU|GPU]"
+}
+
+# check arguments
+if [ $# -ne 1 ]; then
+    usage
+    fail "Error: must specify device"
+fi
+
+if [ "$1" == "CPU" ]; then
+    dev=-1
+elif [ "$1" == "GPU" ]; then
+    export CUDA_VISIBLE_DEVICES=0
+    dev=0
+else
+    usage
+    fail "Unknown device $1"
+fi
+
+pushd $GCN_EXAMPLE_DIR> /dev/null
+
+# test CPU
+python3 gcn.py --dataset cora --gpu $dev || fail "run gcn.py on $1"
+python3 gcn_spmv.py --dataset cora --gpu $dev || fail "run gcn_spmv.py on $1"
+
+popd > /dev/null