remove anonymous repr

examples/pytorch/gcn/gcn_spmv.py

@@ -56,7 +56,7 @@ class GCN(nn.Module):
                 g.apply_nodes(apply_node_func=
                         lambda node: F.dropout(node['h'], p=self.dropout))
             self.g.update_all(fn.copy_src(src='h', out='m'),
-                              fn.sum(msgs='m', out='h'),
+                              fn.sum(msg='m', out='h'),
                               layer)
         return self.g.pop_n_repr('h')
 

python/dgl/__init__.py

@@ -11,5 +11,5 @@ from ._ffi.base import DGLError, __version__
 from .base import ALL
 from .batched_graph import *
 from .generator import *
-from .graph import DGLGraph, __MSG__, __REPR__
+from .graph import DGLGraph
 from .subgraph import DGLSubGraph

python/dgl/base.py

@@ -11,7 +11,4 @@ ALL = "__ALL__"
 def is_all(arg):
     return isinstance(arg, str) and arg == ALL
 
-__MSG__ = "__MSG__"
-__REPR__ = "__REPR__"
-
 dgl_warning = warnings.warn

python/dgl/function/message.py

@@ -4,17 +4,25 @@ from __future__ import absolute_import
 import operator
 import dgl.backend as F
 
-__all__ = ["MessageFunction", "src_mul_edge", "copy_src", "copy_edge"]
+__all__ = ["src_mul_edge", "copy_src", "copy_edge"]
 
 
 class MessageFunction(object):
+    """Base builtin message function class."""
+
     def __call__(self, src, edge):
+        """Regular computation of this builtin.
+
+        This will be used when optimization is not available.
+        """
         raise NotImplementedError
 
     def name(self):
+        """Return the name of this builtin function."""
         raise NotImplementedError
 
     def is_spmv_supported(self, g):
+        """Return whether the SPMV optimization is supported."""
         raise NotImplementedError
 
 
@@ -22,12 +30,6 @@ class BundledMessageFunction(MessageFunction):
     def __init__(self, fn_list):
         if not isinstance(fn_list, (list, tuple)):
             fn_list = [fn_list]
-        else:
-            # sanity check on out field
-            for fn in fn_list:
-                # cannot perform check for udf
-                if isinstance(fn, MessageFunction) and fn.out_field is None:
-                    raise RuntimeError("Not specifying out field for multiple message is ambiguous")
         self.fn_list = fn_list
 
     def is_spmv_supported(self, g):
@@ -43,11 +45,8 @@ class BundledMessageFunction(MessageFunction):
             if ret is None:
                 ret = msg
             else:
-                try:
-                    # ret and msg must be dict
-                    ret.update(msg)
-                except:
-                    raise RuntimeError("Must specify out field for multiple message")
+                # ret and msg must be dict
+                ret.update(msg)
         return ret
 
     def name(self):
@@ -55,25 +54,26 @@ class BundledMessageFunction(MessageFunction):
 
 
 def _is_spmv_supported_node_feat(g, field):
-    if field is None:
-        feat = g.get_n_repr()
-    else:
-        feat = g.get_n_repr()[field]
+    """Return whether the node feature shape supports SPMV optimization.
+
+    Only scalar and vector features are supported currently.
+    """
+    feat = g.get_n_repr()[field]
     shape = F.shape(feat)
     return len(shape) == 1 or len(shape) == 2
 
 def _is_spmv_supported_edge_feat(g, field):
-    # check shape, only scalar edge feature can be optimized at the moment
-    if field is None:
-        feat = g.get_e_repr()
-    else:
-        feat = g.get_e_repr()[field]
+    """Return whether the edge feature shape supports SPMV optimization.
+
+    Only scalar feature is supported currently.
+    """
+    feat = g.get_e_repr()[field]
     shape = F.shape(feat)
     return len(shape) == 1 or (len(shape) == 2 and shape[1] == 1)
 
 
 class SrcMulEdgeMessageFunction(MessageFunction):
-    def __init__(self, mul_op, src_field=None, edge_field=None, out_field=None):
+    def __init__(self, mul_op, src_field, edge_field, out_field):
         self.mul_op = mul_op
         self.src_field = src_field
         self.edge_field = edge_field
@@ -84,21 +84,14 @@ class SrcMulEdgeMessageFunction(MessageFunction):
                 and _is_spmv_supported_edge_feat(g, self.edge_field)
 
     def __call__(self, src, edge):
-        if self.src_field is not None:
-            src = src[self.src_field]
-        if self.edge_field is not None:
-            edge = edge[self.edge_field]
-        ret = self.mul_op(src, edge)
-        if self.out_field is None:
-            return ret
-        else:
-            return {self.out_field : ret}
+        ret = self.mul_op(src[self.src_field], edge[self.edge_field])
+        return {self.out_field : ret}
 
     def name(self):
         return "src_mul_edge"
 
 class CopySrcMessageFunction(MessageFunction):
-    def __init__(self, src_field=None, out_field=None):
+    def __init__(self, src_field, out_field):
         self.src_field = src_field
         self.out_field = out_field
 
@@ -106,14 +99,7 @@ class CopySrcMessageFunction(MessageFunction):
         return _is_spmv_supported_node_feat(g, self.src_field)
 
     def __call__(self, src, edge):
-        if self.src_field is not None:
-            ret = src[self.src_field]
-        else:
-            ret = src
-        if self.out_field is None:
-            return ret
-        else:
-            return {self.out_field : ret}
+        return {self.out_field : src[self.src_field]}
 
     def name(self):
         return "copy_src"
@@ -142,14 +128,41 @@ class CopyEdgeMessageFunction(MessageFunction):
         return "copy_edge"
 
 
-def src_mul_edge(src=None, edge=None, out=None):
-    """TODO(minjie): docstring """
+def src_mul_edge(src, edge, out):
+    """Builtin message function that computes message by multiplying source node features
+    with edge features.
+
+    Parameters
+    ----------
+    src : str
+        The source feature name.
+    edge : str
+        The edge feature name.
+    out : str
+        The output message name.
+    """
     return SrcMulEdgeMessageFunction(operator.mul, src, edge, out)
 
-def copy_src(src=None, out=None):
-    """TODO(minjie): docstring """
+def copy_src(src, out):
+    """Builtin message function that computes message using source node feature.
+
+    Parameters
+    ----------
+    src : str
+        The source feature name.
+    out : str
+        The output message name.
+    """
     return CopySrcMessageFunction(src, out)
 
-def copy_edge(edge=None, out=None):
-    """TODO(minjie): docstring """
+def copy_edge(edge, out):
+    """Builtin message function that computes message using edge feature.
+
+    Parameters
+    ----------
+    edge : str
+        The edge feature name.
+    out : str
+        The output message name.
+    """
     return CopyEdgeMessageFunction(edge, out)

python/dgl/function/reducer.py

@@ -3,27 +3,30 @@ from __future__ import absolute_import
 
 from .. import backend as F
 
-__all__ = ["ReduceFunction", "sum", "max"]
+__all__ = ["sum", "max"]
 
 class ReduceFunction(object):
+    """Base builtin reduce function class."""
+
     def __call__(self, node, msgs):
+        """Regular computation of this builtin.
+
+        This will be used when optimization is not available.
+        """
         raise NotImplementedError
 
     def name(self):
+        """Return the name of this builtin function."""
         raise NotImplementedError
 
     def is_spmv_supported(self):
+        """Return whether the SPMV optimization is supported."""
         raise NotImplementedError
 
 class BundledReduceFunction(ReduceFunction):
     def __init__(self, fn_list):
         if not isinstance(fn_list, (list, tuple)):
             fn_list = [fn_list]
-        else:
-            # sanity check on out field
-            for fn in fn_list:
-                if isinstance(fn, ReduceFunction) and fn.out_field is None:
-                    raise RuntimeError("Not specifying out field for multiple reduce is ambiguous")
         self.fn_list = fn_list
 
     def is_spmv_supported(self):
@@ -39,51 +42,50 @@ class BundledReduceFunction(ReduceFunction):
             if ret is None:
                 ret = rpr
             else:
-                try:
-                    # ret and rpr must be dict
-                    ret.update(rpr)
-                except:
-                    raise RuntimeError("Must specify out field for multiple reudce")
+                # ret and rpr must be dict
+                ret.update(rpr)
         return ret
 
     def name(self):
         return "bundled"
 
 class ReducerFunctionTemplate(ReduceFunction):
-    def __init__(self, name, batch_op, nonbatch_op, msg_field=None, out_field=None):
+    def __init__(self, name, op, msg_field, out_field):
         self.name = name
-        self.batch_op = batch_op
-        self.nonbatch_op = nonbatch_op
+        self.op = op
         self.msg_field = msg_field
         self.out_field = out_field
 
     def is_spmv_supported(self):
-        # TODO: support max
+        # NOTE: only sum is supported right now.
         return self.name == "sum"
 
     def __call__(self, node, msgs):
-        if isinstance(msgs, list):
-            if self.msg_field is None:
-                ret = self.nonbatch_op(msgs)
-            else:
-                ret = self.nonbatch_op([msg[self.msg_field] for msg in msgs])
-        else:
-            if self.msg_field is None:
-                ret = self.batch_op(msgs, 1)
-            else:
-                ret = self.batch_op(msgs[self.msg_field], 1)
-        if self.out_field is None:
-            return ret
-        else:
-            return {self.out_field : ret}
+        return {self.out_field : self.op(msgs[self.msg_field], 1)}
 
     def name(self):
         return self.name
 
-_python_sum = sum
-def sum(msgs=None, out=None):
-    return ReducerFunctionTemplate("sum", F.sum, _python_sum, msgs, out)
+def sum(msg, out):
+    """Builtin reduce function that aggregates messages by sum.
+
+    Parameters
+    ----------
+    msg : str
+        The message name.
+    out : str
+        The output node feature name.
+    """
+    return ReducerFunctionTemplate("sum", F.sum, msg, out)
+
+def max(msg, out):
+    """Builtin reduce function that aggregates messages by max.
 
-_python_max = max
-def max(msgs=None, out=None):
-    return ReducerFunctionTemplate("max", F.max, _python_max, msgs, out)
+    Parameters
+    ----------
+    msg : str
+        The message name.
+    out : str
+        The output node feature name.
+    """
+    return ReducerFunctionTemplate("max", F.max, msg, out)

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
+        for key, val in hu.items():
+            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
+            for key, val in hu.items():
+                self._node_frame[key] = val
         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)
+            self._node_frame.update_rows(u, 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)
+            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)
+            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():
-                    self._edge_frame[key] = val
-            else:
-                self._edge_frame[__REPR__] = h_uv
+            for key, val in he.items():
+                self._edge_frame[key] = val
         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,21 +759,15 @@ 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)
+            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)
+            return self._edge_frame.select_rows(eid)
 
     def register_edge_func(self, edge_func):
         """Register global edge update function.
@@ -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})
+        self._msg_frame.append(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)
+            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)}
+        keys = new_reprs[0].keys()
+        new_reprs = {key : F.pack([repr[key] for repr in new_reprs])
+                     for key in keys}
 
         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})
+            msg_frame.append(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,10 +134,7 @@ 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]
+        srccol = self.node_repr[self.src_field]
         ctx = F.get_context(srccol)
 
         # build adjmat
@@ -142,10 +147,7 @@ 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}
+        return {self.dst_field : dstcol}
 
 
 # FIXME: refactorize in scheduler/executor redesign
@@ -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)
+            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)}
+        keys = new_reprs[0].keys()
+        new_reprs = {key : F.pack([repr[key] for repr in new_reprs])
+                     for key in keys}
         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 = 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,10 +338,7 @@ 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 = self.edge_repr[edge_field]
             dat = F.squeeze(dat)
         else:
             dat = F.ones((len(self.u), ))
@@ -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)

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():

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)
@@ -127,19 +125,19 @@ def test_update_all_multi_fn():
     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)
@@ -147,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']
@@ -181,20 +179,23 @@ def test_send_and_recv_multi_fn():
     # 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)
@@ -203,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']