Merge branch 'cpp' of github.com:jermainewang/dgl into cpp

python/dgl/graph.py

@@ -6,7 +6,7 @@ import networkx as nx
 import numpy as np
 
 import dgl
-from .base import ALL, is_all, __MSG__, __REPR__
+from .base import ALL, is_all, DGLError, dgl_warning
 from . import backend as F
 from .backend import Tensor
 from .frame import FrameRef, merge_frames
@@ -22,7 +22,6 @@ class DGLGraph(object):
     """Base graph class specialized for neural networks on graphs.
 
     TODO(minjie): document of batching semantics
-    TODO(minjie): document of __REPR__ semantics
 
     Parameters
     ----------
@@ -448,7 +447,9 @@ class DGLGraph(object):
             The nx graph
         """
         nx_graph = self._graph.to_networkx()
-        #TODO: attributes
+        #TODO(minjie): attributes
+        dgl_warning('to_networkx currently does not support converting'
+                    ' node/edge features automatically.')
         return nx_graph
 
     def from_networkx(self, nx_graph, node_attrs=None, edge_attrs=None):
@@ -550,20 +551,17 @@ class DGLGraph(object):
     def set_n_repr(self, hu, u=ALL, inplace=False):
         """Set node(s) representation.
 
-        To set multiple node representations at once, pass `u` with a tensor or
-        a supported container of node ids. In this case, `hu` must be a tensor
-        of shape (B, D1, D2, ...), where B is the number of the nodes and
-        (D1, D2, ...) is the shape of the node representation tensor.
-
-        Dictionary type is also supported for `hu`. In this case, each item
-        will be treated as separate attribute of the nodes.
+        `hu` is a dictionary from the feature name to feature tensor. Each tensor
+        is of shape (B, D1, D2, ...), where B is the number of nodes to be updated,
+        and (D1, D2, ...) be the shape of the node representation tensor. The
+        length of the given node ids must match B (i.e, len(u) == B).
 
         All update will be done out-placely to work with autograd unless the inplace
         flag is true.
 
         Parameters
         ----------
-        hu : tensor or dict of tensor
+        hu : dict of tensor
             Node representation.
         u : node, container or tensor
             The node(s).
@@ -571,32 +569,31 @@ class DGLGraph(object):
             True if the update is done inplacely
         """
         # sanity check
+        if not utils.is_dict_like(hu):
+            raise DGLError('Expect dictionary type for feature data.'
+                           ' Got "%s" instead.' % type(hu))
         if is_all(u):
             num_nodes = self.number_of_nodes()
         else:
             u = utils.toindex(u)
             num_nodes = len(u)
-        if utils.is_dict_like(hu):
         for key, val in hu.items():
-                assert F.shape(val)[0] == num_nodes
-        else:
-            assert F.shape(hu)[0] == num_nodes
+            nfeats = F.shape(val)[0]
+            if nfeats != num_nodes:
+                raise DGLError('Expect number of features to match number of nodes (len(u)).'
+                               ' Got %d and %d instead.' % (nfeats, num_nodes))
         # set
         if is_all(u):
-            if utils.is_dict_like(hu):
             for key, val in hu.items():
                 self._node_frame[key] = val
         else:
-                self._node_frame[__REPR__] = hu
-        else:
-            if utils.is_dict_like(hu):
             self._node_frame.update_rows(u, hu, inplace=inplace)
-            else:
-                self._node_frame.update_rows(u, {__REPR__ : hu}, inplace=inplace)
 
     def get_n_repr(self, u=ALL):
         """Get node(s) representation.
 
+        The returned feature tensor batches multiple node features on the first dimension.
+
         Parameters
         ----------
         u : node, container or tensor
@@ -605,23 +602,17 @@ class DGLGraph(object):
         Returns
         -------
         dict
-            Representation dict
+            Representation dict from feature name to feature tensor.
         """
         if len(self.node_attr_schemes()) == 0:
             return dict()
         if is_all(u):
-            if len(self._node_frame) == 1 and __REPR__ in self._node_frame:
-                return self._node_frame[__REPR__]
-            else:
             return dict(self._node_frame)
         else:
             u = utils.toindex(u)
-            if len(self._node_frame) == 1 and __REPR__ in self._node_frame:
-                return self._node_frame.select_rows(u)[__REPR__]
-            else:
             return self._node_frame.select_rows(u)
 
-    def pop_n_repr(self, key=__REPR__):
+    def pop_n_repr(self, key):
         """Get and remove the specified node repr.
 
         Parameters
@@ -636,23 +627,19 @@ class DGLGraph(object):
         """
         return self._node_frame.pop(key)
 
-    def set_e_repr(self, h_uv, u=ALL, v=ALL, inplace=False):
+    def set_e_repr(self, he, u=ALL, v=ALL, inplace=False):
         """Set edge(s) representation.
 
-        To set multiple edge representations at once, pass `u` and `v` with tensors or
-        supported containers of node ids. In this case, `h_uv` must be a tensor
-        of shape (B, D1, D2, ...), where B is the number of the edges and
-        (D1, D2, ...) is the shape of the edge representation tensor.
-
-        Dictionary type is also supported for `h_uv`. In this case, each item
-        will be treated as separate attribute of the edges.
+        `he` is a dictionary from the feature name to feature tensor. Each tensor
+        is of shape (B, D1, D2, ...), where B is the number of edges to be updated,
+        and (D1, D2, ...) be the shape of the edge representation tensor.
 
         All update will be done out-placely to work with autograd unless the inplace
         flag is true.
 
         Parameters
         ----------
-        h_uv : tensor or dict of tensor
+        he : tensor or dict of tensor
           Edge representation.
         u : node, container or tensor
           The source node(s).
@@ -662,26 +649,33 @@ class DGLGraph(object):
             True if the update is done inplacely
         """
         # sanity check
+        if not utils.is_dict_like(he):
+            raise DGLError('Expect dictionary type for feature data.'
+                           ' Got "%s" instead.' % type(he))
         u_is_all = is_all(u)
         v_is_all = is_all(v)
         assert u_is_all == v_is_all
         if u_is_all:
-            self.set_e_repr_by_id(h_uv, eid=ALL, inplace=inplace)
+            self.set_e_repr_by_id(he, eid=ALL, inplace=inplace)
         else:
             u = utils.toindex(u)
             v = utils.toindex(v)
             _, _, eid = self._graph.edge_ids(u, v)
-            self.set_e_repr_by_id(h_uv, eid=eid, inplace=inplace)
+            self.set_e_repr_by_id(he, eid=eid, inplace=inplace)
 
-    def set_e_repr_by_id(self, h_uv, eid=ALL, inplace=False):
+    def set_e_repr_by_id(self, he, eid=ALL, inplace=False):
         """Set edge(s) representation by edge id.
 
+        `he` is a dictionary from the feature name to feature tensor. Each tensor
+        is of shape (B, D1, D2, ...), where B is the number of edges to be updated,
+        and (D1, D2, ...) be the shape of the edge representation tensor.
+
         All update will be done out-placely to work with autograd unless the inplace
         flag is true.
 
         Parameters
         ----------
-        h_uv : tensor or dict of tensor
+        he : tensor or dict of tensor
           Edge representation.
         eid : int, container or tensor
           The edge id(s).
@@ -689,30 +683,27 @@ class DGLGraph(object):
             True if the update is done inplacely
         """
         # sanity check
+        if not utils.is_dict_like(he):
+            raise DGLError('Expect dictionary type for feature data.'
+                           ' Got "%s" instead.' % type(he))
         if is_all(eid):
             num_edges = self.number_of_edges()
         else:
             eid = utils.toindex(eid)
             num_edges = len(eid)
-        if utils.is_dict_like(h_uv):
-            for key, val in h_uv.items():
-                assert F.shape(val)[0] == num_edges
-        else:
-            assert F.shape(h_uv)[0] == num_edges
+        for key, val in he.items():
+            nfeats = F.shape(val)[0]
+            if nfeats != num_edges:
+                raise DGLError('Expect number of features to match number of edges.'
+                               ' Got %d and %d instead.' % (nfeats, num_edges))
         # set
         if is_all(eid):
             # update column
-            if utils.is_dict_like(h_uv):
-                for key, val in h_uv.items():
+            for key, val in he.items():
                 self._edge_frame[key] = val
-            else:
-                self._edge_frame[__REPR__] = h_uv
         else:
             # update row
-            if utils.is_dict_like(h_uv):
-                self._edge_frame.update_rows(eid, h_uv, inplace=inplace)
-            else:
-                self._edge_frame.update_rows(eid, {__REPR__ : h_uv}, inplace=inplace)
+            self._edge_frame.update_rows(eid, he, inplace=inplace)
 
     def get_e_repr(self, u=ALL, v=ALL):
         """Get node(s) representation.
@@ -742,7 +733,7 @@ class DGLGraph(object):
             _, _, eid = self._graph.edge_ids(u, v)
             return self.get_e_repr_by_id(eid=eid)
 
-    def pop_e_repr(self, key=__REPR__):
+    def pop_e_repr(self, key):
         """Get and remove the specified edge repr.
 
         Parameters
@@ -768,20 +759,14 @@ class DGLGraph(object):
         Returns
         -------
         dict
-            Representation dict
+            Representation dict from feature name to feature tensor.
         """
         if len(self.edge_attr_schemes()) == 0:
             return dict()
         if is_all(eid):
-            if len(self._edge_frame) == 1 and __REPR__ in self._edge_frame:
-                return self._edge_frame[__REPR__]
-            else:
             return dict(self._edge_frame)
         else:
             eid = utils.toindex(eid)
-            if len(self._edge_frame) == 1 and __REPR__ in self._edge_frame:
-                return self._edge_frame.select_rows(eid)[__REPR__]
-            else:
             return self._edge_frame.select_rows(eid)
 
     def register_edge_func(self, edge_func):
@@ -837,6 +822,8 @@ class DGLGraph(object):
     def apply_nodes(self, v=ALL, apply_node_func="default"):
         """Apply the function on node representations.
 
+        Applying a None function will be ignored.
+
         Parameters
         ----------
         v : int, iterable of int, tensor, optional
@@ -868,7 +855,7 @@ class DGLGraph(object):
             # merge current node_repr with reduce output
             curr_repr = utils.HybridDict(reduce_accum, curr_repr)
         new_repr = apply_node_func(curr_repr)
-        if reduce_accum is not None and utils.is_dict_like(new_repr) :
+        if reduce_accum is not None:
             # merge new node_repr with reduce output
             reduce_accum.update(new_repr)
             new_repr = reduce_accum
@@ -877,6 +864,8 @@ class DGLGraph(object):
     def apply_edges(self, u=None, v=None, apply_edge_func="default", eid=None):
         """Apply the function on edge representations.
 
+        Applying a None function will be ignored.
+
         Parameters
         ----------
         u : optional, int, iterable of int, tensor
@@ -893,7 +882,6 @@ class DGLGraph(object):
         if not apply_edge_func:
             # Skip none function call.
             return
-
         if eid is None:
             new_repr = apply_edge_func(self.get_e_repr(u, v))
             self.set_e_repr(new_repr, u, v)
@@ -914,9 +902,8 @@ class DGLGraph(object):
         The message function can be any of the pre-defined functions
         ('from_src').
 
-        Currently, we require the message functions of consecutive send's and
-        send_on's to return the same keys.  Otherwise the behavior will be
-        undefined.
+        Currently, we require the message functions of consecutive send's to
+        return the same keys.  Otherwise the behavior will be undefined.
 
         Parameters
         ----------
@@ -964,10 +951,7 @@ class DGLGraph(object):
             edge_reprs = self.get_e_repr_by_id(eid)
             msgs = message_func(src_reprs, edge_reprs)
         self._msg_graph.add_edges(u, v)
-        if utils.is_dict_like(msgs):
         self._msg_frame.append(msgs)
-        else:
-            self._msg_frame.append({__MSG__ : msgs})
 
         # TODO(minjie): Fix these codes in next PR.
         """
@@ -1061,7 +1045,6 @@ class DGLGraph(object):
                 v = utils.toindex(v)
                 u, v = utils.edge_broadcasting(u, v)
                 _, _, eid = self._graph.edge_ids(u, v)
-
             # call the UDF
             src_reprs = self.get_n_repr(u)
             dst_reprs = self.get_n_repr(v)
@@ -1148,25 +1131,19 @@ class DGLGraph(object):
                 msg_shape = F.shape(msg)
                 new_shape = (bkt_len, deg) + msg_shape[1:]
                 return F.reshape(msg, new_shape)
-            if len(in_msgs) == 1 and __MSG__ in in_msgs:
-                reshaped_in_msgs = _reshape_fn(in_msgs[__MSG__])
-            else:
             reshaped_in_msgs = utils.LazyDict(
                     lambda key: _reshape_fn(in_msgs[key]), self._msg_frame.schemes)
             reordered_v.append(v_bkt.tousertensor())
             new_reprs.append(reduce_func(dst_reprs, reshaped_in_msgs))
 
-        # TODO: clear partial messages
+        # TODO(minjie): clear partial messages
         self.reset_messages()
 
         # Pack all reducer results together
         reordered_v = F.pack(reordered_v)
-        if utils.is_dict_like(new_reprs[0]):
         keys = new_reprs[0].keys()
         new_reprs = {key : F.pack([repr[key] for repr in new_reprs])
                      for key in keys}
-        else:
-            new_reprs = {__REPR__ : F.pack(new_reprs)}
 
         if v_is_all and not has_zero_degree:
             # First do reorder and then replace the whole column.
@@ -1237,15 +1214,13 @@ class DGLGraph(object):
 
         if executor:
             new_reprs = executor.run()
-            if not utils.is_dict_like(new_reprs):
-                new_reprs = {__REPR__: new_reprs}
             unique_v = executor.recv_nodes
             self._apply_nodes(unique_v, apply_node_func, reduce_accum=new_reprs)
         elif eid is not None:
             _, v, _ = self._graph.find_edges(eid)
             unique_v = utils.toindex(F.unique(v.tousertensor()))
 
-            # TODO: replace with the new DegreeBucketingScheduler
+            # TODO(quan): replace with the new DegreeBucketingScheduler
             self.send(eid=eid, message_func=message_func)
             self.recv(unique_v, reduce_func, apply_node_func)
         else:
@@ -1261,10 +1236,7 @@ class DGLGraph(object):
             edge_reprs = self.get_e_repr(u, v)
             msgs = message_func(src_reprs, edge_reprs)
             msg_frame = FrameRef()
-            if utils.is_dict_like(msgs):
             msg_frame.append(msgs)
-            else:
-                msg_frame.append({__MSG__: msgs})
 
             # recv with degree bucketing
             executor = scheduler.get_recv_executor(graph=self,
@@ -1353,8 +1325,6 @@ class DGLGraph(object):
                 "update_all", self, message_func=message_func, reduce_func=reduce_func)
         if executor:
             new_reprs = executor.run()
-            if not utils.is_dict_like(new_reprs):
-                new_reprs = {__REPR__: new_reprs}
             self._apply_nodes(ALL, apply_node_func, reduce_accum=new_reprs)
         else:
             self.send(ALL, ALL, message_func)
@@ -1387,7 +1357,7 @@ class DGLGraph(object):
             Arguments for pre-defined iterators.
         """
         if isinstance(traverser, str):
-            # TODO Call pre-defined routine to unroll the computation.
+            # TODO(minjie): Call pre-defined routine to unroll the computation.
             raise RuntimeError('Not implemented.')
         else:
             # NOTE: the iteration can return multiple edges at each step.

python/dgl/scheduler.py

@@ -3,7 +3,7 @@ from __future__ import absolute_import
 
 import numpy as np
 
-from .base import ALL, __MSG__, __REPR__
+from .base import ALL, DGLError
 from . import backend as F
 from .function import message as fmsg
 from .function import reducer as fred
@@ -111,7 +111,15 @@ def light_degree_bucketing_for_graph(graph):
 
 
 class Executor(object):
+    """Base class for executing graph computation."""
+
     def run(self):
+        """Run this executor.
+
+        This should return the new node features.
+
+        TODO(minjie): extend this to support computation on edges.
+        """
         raise NotImplementedError
 
 class SPMVOperator(Executor):
@@ -126,9 +134,6 @@ class SPMVOperator(Executor):
 
     def run(self):
         # get src col
-        if self.src_field is None:
-            srccol = self.node_repr
-        else:
         srccol = self.node_repr[self.src_field]
         ctx = F.get_context(srccol)
 
@@ -142,9 +147,6 @@ class SPMVOperator(Executor):
             dstcol = F.squeeze(dstcol)
         else:
             dstcol = F.spmm(adjmat, srccol)
-        if self.dst_field is None:
-            return dstcol
-        else:
         return {self.dst_field : dstcol}
 
 
@@ -180,20 +182,14 @@ class DegreeBucketingExecutor(Executor):
                 msg_shape = F.shape(msg)
                 new_shape = (len(vv), deg) + msg_shape[1:]
                 return F.reshape(msg, new_shape)
-            if len(in_msgs) == 1 and __MSG__ in in_msgs:
-                reshaped_in_msgs = _reshape_fn(in_msgs[__MSG__])
-            else:
             reshaped_in_msgs = utils.LazyDict(
                     lambda key: _reshape_fn(in_msgs[key]), self.msg_frame.schemes)
             new_reprs.append(self.rfunc(dst_reprs, reshaped_in_msgs))
 
         # Pack all reducer results together
-        if utils.is_dict_like(new_reprs[0]):
         keys = new_reprs[0].keys()
         new_reprs = {key : F.pack([repr[key] for repr in new_reprs])
                      for key in keys}
-        else:
-            new_reprs = {__REPR__ : F.pack(new_reprs)}
         return new_reprs
 
 
@@ -249,12 +245,6 @@ class UpdateAllExecutor(BasicExecutor):
         self._graph_shape = None
         self._recv_nodes = None
 
-    @property
-    def graph_idx(self):
-        if self._graph_idx is None:
-            self._graph_idx = self.g._graph.adjacency_matrix()
-        return self._graph_idx
-
     @property
     def graph_shape(self):
         if self._graph_shape is None:
@@ -280,16 +270,13 @@ class UpdateAllExecutor(BasicExecutor):
 
     def _adj_build_fn(self, edge_field, ctx, use_edge_feat):
         if use_edge_feat:
-            if edge_field is None:
-                dat = self.edge_repr
-            else:
             dat = self.edge_repr[edge_field]
             dat = F.squeeze(dat)
             # TODO(minjie): should not directly use _indices
-            idx = self.graph_idx.get(ctx)._indices()
+            idx = self.g.adjacency_matrix(ctx)._indices()
             adjmat = F.sparse_tensor(idx, dat, self.graph_shape)
         else:
-            adjmat = self.graph_idx.get(ctx)
+            adjmat = self.g.adjacency_matrix(ctx)
         return adjmat
 
 
@@ -351,9 +338,6 @@ class SendRecvExecutor(BasicExecutor):
 
     def _adj_build_fn(self, edge_field, ctx, use_edge_feat):
         if use_edge_feat:
-            if edge_field is None:
-                dat = self.edge_repr
-            else:
             dat = self.edge_repr[edge_field]
             dat = F.squeeze(dat)
         else:
@@ -386,9 +370,8 @@ class BundledExecutor(BasicExecutor):
         func_pairs = []
         for rfn in rfunc.fn_list:
             mfn = out2mfunc.get(rfn.msg_field, None)
-            # field check
-            assert mfn is not None, \
-                    "cannot find message func for reduce func in-field {}".format(rfn.msg_field)
+            if mfn is None:
+                raise DGLError('Cannot find message field "%s".' % rfn.msg_field)
             func_pairs.append((mfn, rfn))
         return func_pairs
 
@@ -409,7 +392,6 @@ class BundledUpdateAllExecutor(BundledExecutor, UpdateAllExecutor):
         self._init_state()
         BundledExecutor.__init__(self, graph, mfunc, rfunc)
 
-
 class BundledSendRecvExecutor(BundledExecutor, SendRecvExecutor):
     def __init__(self, graph, src, dst, mfunc, rfunc):
         self._init_state(src, dst)

src/c_api_common.cc

+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file c_runtime_api.cc
+ * \brief DGL C API common implementations
+ */
 #include "c_api_common.h"
 
 using tvm::runtime::TVMArgs;

src/c_api_common.h

-// DGL C API common util functions
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file c_api_common.h
+ * \brief DGL C API common util functions
+ */
 #ifndef DGL_C_API_COMMON_H_
 #define DGL_C_API_COMMON_H_
 
@@ -12,11 +16,19 @@ namespace dgl {
 // Graph handler type
 typedef void* GraphHandle;
 
-// Convert the given DLTensor to a temporary DLManagedTensor that does not own memory.
-DLManagedTensor* CreateTmpDLManagedTensor(const tvm::runtime::TVMArgValue& arg);
+/*!
+ * \brief Convert the given DLTensor to DLManagedTensor.
+ *
+ * Return a temporary DLManagedTensor that does not own memory.
+ */
+DLManagedTensor* CreateTmpDLManagedTensor(
+    const tvm::runtime::TVMArgValue& arg);
 
-// Convert a vector of NDArray to PackedFunc
-tvm::runtime::PackedFunc ConvertNDArrayVectorToPackedFunc(const std::vector<tvm::runtime::NDArray>& vec);
+/*!
+ * \brief Convert a vector of NDArray to PackedFunc.
+ */
+tvm::runtime::PackedFunc ConvertNDArrayVectorToPackedFunc(
+    const std::vector<tvm::runtime::NDArray>& vec);
 
 }  // namespace dgl
 

src/graph/graph.cc

-// Graph class implementation
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file graph/graph.cc
+ * \brief DGL graph index implementation
+ */
+#include <dgl/graph.h>
 #include <algorithm>
 #include <unordered_map>
 #include <set>
 #include <functional>
-#include <dgl/graph.h>
 
 namespace dgl {
 namespace {
@@ -461,7 +465,8 @@ Subgraph Graph::EdgeSubgraph(IdArray eids) const {
     rst.graph.AddEdge(oldv2newv[src_id], oldv2newv[dst_id]);
   }
 
-  rst.induced_vertices = IdArray::Empty({static_cast<int64_t>(nodes.size())}, eids->dtype, eids->ctx);
+  rst.induced_vertices = IdArray::Empty(
+      {static_cast<int64_t>(nodes.size())}, eids->dtype, eids->ctx);
   std::copy(nodes.begin(), nodes.end(), static_cast<int64_t*>(rst.induced_vertices->data));
 
   return rst;

src/graph/graph_apis.cc

+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file graph/graph.cc
+ * \brief DGL graph index APIs
+ */
 #include <dgl/graph.h>
 #include <dgl/graph_op.h>
 #include "../c_api_common.h"

src/graph/graph_op.cc

-// Graph operation implementation
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file graph/graph.cc
+ * \brief Graph operation implementation
+ */
 #include <dgl/graph_op.h>
 #include <algorithm>
 
 namespace dgl {
 
-Graph GraphOp::LineGraph(const Graph* g, bool backtracking){
+Graph GraphOp::LineGraph(const Graph* g, bool backtracking) {
   typedef std::pair<dgl_id_t, dgl_id_t> entry;
   typedef std::map<dgl_id_t, std::vector<entry>> csm;  // Compressed Sparse Matrix
 
@@ -117,32 +121,6 @@ std::vector<Graph> GraphOp::DisjointPartitionBySizes(const Graph* graph, IdArray
     node_offset += sizes_data[i];
     edge_offset += num_edges;
   }
-  /*for (int64_t i = 0; i < len; ++i) {
-    rst[i].AddVertices(sizes_data[i]);
-  }
-  for (dgl_id_t eid = 0; eid < graph->num_edges_; ++eid) {
-    const dgl_id_t src = graph->all_edges_src_[eid];
-    const dgl_id_t dst = graph->all_edges_dst_[eid];
-    size_t src_select = 0, dst_select = 0;
-    for (size_t i = 1; i < cumsum.size(); ++i) { // TODO: replace with binary search
-      if (cumsum[i] > src) {
-        src_select = i;
-        break;
-      }
-    }
-    for (size_t i = 1; i < cumsum.size(); ++i) { // TODO: replace with binary search
-      if (cumsum[i] > dst) {
-        dst_select = i;
-        break;
-      }
-    }
-    if (src_select != dst_select) {
-      // the edge is ignored if across two partitions
-      continue;
-    }
-    const int64_t offset = cumsum[src_select - 1];
-    rst[src_select - 1].AddEdge(src - offset, dst - offset);
-  }*/
   return rst;
 }
 

src/runtime/README.md

-# C API and runtime
-
-Borrowed and adapted from TVM project.

src/runtime/file_util.h

@@ -3,8 +3,8 @@
  * \file file_util.h
  * \brief Minimum file manipulation util for runtime.
  */
-#ifndef TVM_RUNTIME_FILE_UTIL_H_
-#define TVM_RUNTIME_FILE_UTIL_H_
+#ifndef DGL_RUNTIME_FILE_UTIL_H_
+#define DGL_RUNTIME_FILE_UTIL_H_
 
 #include <string>
 #include "meta_data.h"
@@ -73,4 +73,4 @@ void LoadMetaDataFromFile(
     std::unordered_map<std::string, FunctionInfo>* fmap);
 }  // namespace runtime
 }  // namespace tvm
-#endif  // TVM_RUNTIME_FILE_UTIL_H_
+#endif  // DGL_RUNTIME_FILE_UTIL_H_

src/runtime/meta_data.h

@@ -3,8 +3,8 @@
  * \file meta_data.h
  * \brief Meta data related utilities
  */
-#ifndef TVM_RUNTIME_META_DATA_H_
-#define TVM_RUNTIME_META_DATA_H_
+#ifndef DGL_RUNTIME_META_DATA_H_
+#define DGL_RUNTIME_META_DATA_H_
 
 #include <dmlc/json.h>
 #include <dmlc/io.h>
@@ -33,4 +33,4 @@ struct FunctionInfo {
 namespace dmlc {
 DMLC_DECLARE_TRAITS(has_saveload, ::tvm::runtime::FunctionInfo, true);
 }  // namespace dmlc
-#endif  // TVM_RUNTIME_META_DATA_H_
+#endif  // DGL_RUNTIME_META_DATA_H_

src/runtime/module_util.h

@@ -3,8 +3,8 @@
  * \file module_util.h
  * \brief Helper utilities for module building
  */
-#ifndef TVM_RUNTIME_MODULE_UTIL_H_
-#define TVM_RUNTIME_MODULE_UTIL_H_
+#ifndef DGL_RUNTIME_MODULE_UTIL_H_
+#define DGL_RUNTIME_MODULE_UTIL_H_
 
 #include <dgl/runtime/module.h>
 #include <dgl/runtime/c_runtime_api.h>
@@ -58,4 +58,4 @@ void InitContextFunctions(FLookup flookup) {
 }
 }  // namespace runtime
 }  // namespace tvm
-#endif   // TVM_RUNTIME_MODULE_UTIL_H_
+#endif   // DGL_RUNTIME_MODULE_UTIL_H_

src/runtime/pack_args.h

@@ -10,8 +10,8 @@
  *                 union_32bit args[N], int num_args);
  *      - Pack buffer by address, pack rest parameter into 32bit union buffer.
  */
-#ifndef TVM_RUNTIME_PACK_ARGS_H_
-#define TVM_RUNTIME_PACK_ARGS_H_
+#ifndef DGL_RUNTIME_PACK_ARGS_H_
+#define DGL_RUNTIME_PACK_ARGS_H_
 
 #include <dgl/runtime/c_runtime_api.h>
 #include <vector>
@@ -307,4 +307,4 @@ inline PackedFunc PackFuncPackedArg(F f, const std::vector<TVMType>& arg_types)
 }
 }  // namespace runtime
 }  // namespace tvm
-#endif  // TVM_RUNTIME_PACK_ARGS_H_
+#endif  // DGL_RUNTIME_PACK_ARGS_H_

src/runtime/runtime_base.h

@@ -3,8 +3,8 @@
  * \file runtime_base.h
  * \brief Base of all C APIs
  */
-#ifndef TVM_RUNTIME_RUNTIME_BASE_H_
-#define TVM_RUNTIME_RUNTIME_BASE_H_
+#ifndef DGL_RUNTIME_RUNTIME_BASE_H_
+#define DGL_RUNTIME_RUNTIME_BASE_H_
 
 #include <dgl/runtime/c_runtime_api.h>
 #include <stdexcept>
@@ -31,4 +31,4 @@ inline int TVMAPIHandleException(const std::runtime_error &e) {
   return -1;
 }
 
-#endif  // TVM_RUNTIME_RUNTIME_BASE_H_
+#endif  // DGL_RUNTIME_RUNTIME_BASE_H_

src/runtime/thread_storage_scope.h

@@ -3,8 +3,8 @@
  * \file thread_storage_scope.h
  * \brief Extract thread axis configuration from TVMArgs.
  */
-#ifndef TVM_RUNTIME_THREAD_STORAGE_SCOPE_H_
-#define TVM_RUNTIME_THREAD_STORAGE_SCOPE_H_
+#ifndef DGL_RUNTIME_THREAD_STORAGE_SCOPE_H_
+#define DGL_RUNTIME_THREAD_STORAGE_SCOPE_H_
 
 #include <dgl/runtime/packed_func.h>
 #include <string>
@@ -204,4 +204,4 @@ struct hash<::tvm::runtime::StorageScope> {
   }
 };
 }  // namespace std
-#endif  // TVM_RUNTIME_THREAD_STORAGE_SCOPE_H_
+#endif  // DGL_RUNTIME_THREAD_STORAGE_SCOPE_H_

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;
             }
         }

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,7 +23,7 @@ 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));
   });
 

tests/pytorch/test_basics.py

@@ -209,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))
@@ -226,25 +225,25 @@ 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])
 

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)
-    ncol = Variable(th.randn(10, D), requires_grad=grad)
-    ecol = Variable(th.randn(17, D), requires_grad=grad)
-    g.set_n_repr(ncol)
-    g.set_e_repr(ecol)
-    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,23 +123,23 @@ 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():