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_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/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