Repository: dmlc/dgl

Issue:
[BUG] Consecutive sends without recv
Here's the test I was running on `cpp` branch.
```python
def test_send_twice():
    g = DGLGraph()
    g.add_nodes(2)
    g.add_edge(0, 1)
    def _message_a(src, edge):
        return {'a': src['a']}
    def _message_b(src, edge):
        return {'b': src['b']}
    def _reduce(node, msgs):
        return {'a': msgs['a'].sum(1), 'b': msgs['b'].sum(1)}

    old_a = th.randn(2, 5)
    old_b = th.randn(2, 5)
    g.set_n_repr({'a': old_a, 'b': old_b})

    g.send(0, 1, _message_a)
    g.send(0, 1, _message_b)
    g.recv([1], _reduce)

    new_repr = g.get_n_repr()

    assert th.allclose(new_repr['a'][1], old_a[0])
    assert th.allclose(new_repr['b'][1], old_b[0])
```

First, if I do two consecutive sends with different message keys, DGLGraph throws a `KeyError: 'b'` regardless of whether we are sending on the same node pairs/edges.

The next issue is that we are adding edges into a "message graph" during sends.  This does not make a lot of sense on multigraphs if we allow sending between the same node pairs twice without recv, which will leave us with two edges in the message graph.  Actually, in the test case above, I stepped into the `_batch_recv` method and found that the degree is 2 in the degree bucketing procedure.

So do we want to allow (1) consecutive sends without recv, and (2) consecutive sends between the same node pairs without recv?

If we were to allow both, I guess ultimately we would want to have the message graph represented as an edge subgraph, or somehow maintain a set of "already-added-edges" in the message graph.

Thoughts?


Discussion:

Resolving patch:
diff --git a/python/dgl/graph.py b/python/dgl/graph.py
--- a/python/dgl/graph.py
+++ b/python/dgl/graph.py
@@ -32,19 +32,23 @@ class DGLGraph(object):
         Node feature storage.
     edge_frame : FrameRef
         Edge feature storage.
+    multigraph : bool, optional
+        Whether the graph would be a multigraph (default: False)
     """
     def __init__(self,
                  graph_data=None,
                  node_frame=None,
-                 edge_frame=None):
+                 edge_frame=None,
+                 multigraph=False):
         # graph
-        self._graph = create_graph_index(graph_data)
+        self._graph = create_graph_index(graph_data, multigraph)
         # frame
         self._node_frame = node_frame if node_frame is not None else FrameRef()
         self._edge_frame = edge_frame if edge_frame is not None else FrameRef()
         # msg graph & frame
-        self._msg_graph = create_graph_index()
+        self._msg_graph = create_graph_index(multigraph=multigraph)
         self._msg_frame = FrameRef()
+        self._msg_edges = []
         self.reset_messages()
         # registered functions
         self._message_func = None
@@ -109,11 +113,13 @@ def clear(self):
         self._edge_frame.clear()
         self._msg_graph.clear()
         self._msg_frame.clear()
+        self._msg_edges.clear()
 
     def reset_messages(self):
         """Clear all messages."""
         self._msg_graph.clear()
         self._msg_frame.clear()
+        self._msg_edges.clear()
         self._msg_graph.add_nodes(self.number_of_nodes())
 
     def number_of_nodes(self):
@@ -130,6 +136,12 @@ def __len__(self):
         """Return the number of nodes."""
         return self.number_of_nodes()
 
+    @property
+    def is_multigraph(self):
+        """Whether the graph is a multigraph.
+        """
+        return self._graph.is_multigraph()
+
     def number_of_edges(self):
         """Return the number of edges.
 
@@ -176,7 +188,7 @@ def has_nodes(self, vids):
         rst = self._graph.has_nodes(vids)
         return rst.tousertensor()
 
-    def has_edge(self, u, v):
+    def has_edge_between(self, u, v):
         """Return true if the edge exists.
 
         Parameters
@@ -191,9 +203,9 @@ def has_edge(self, u, v):
         bool
             True if the edge exists
         """
-        return self._graph.has_edge(u, v)
+        return self._graph.has_edge_between(u, v)
 
-    def has_edges(self, u, v):
+    def has_edges_between(self, u, v):
         """Return true if the edge exists.
 
         Parameters
@@ -210,7 +222,7 @@ def has_edges(self, u, v):
         """
         u = utils.toindex(u)
         v = utils.toindex(v)
-        rst = self._graph.has_edges(u, v)
+        rst = self._graph.has_edges_between(u, v)
         return rst.tousertensor()
 
     def predecessors(self, v, radius=1):
@@ -247,7 +259,7 @@ def successors(self, v, radius=1):
         """
         return self._graph.successors(v).tousertensor()
 
-    def edge_id(self, u, v):
+    def edge_id(self, u, v, force_multi=False):
         """Return the id of the edge.
 
         Parameters
@@ -256,15 +268,20 @@ def edge_id(self, u, v):
             The src node.
         v : int
             The dst node.
+        force_multi : bool
+            If False, will return a single edge ID if the graph is a simple graph.
+            If True, will always return an array.
 
         Returns
         -------
-        int
-            The edge id.
+        int or tensor
+            The edge id if force_multi == True and the graph is a simple graph.
+            The edge id array otherwise.
         """
-        return self._graph.edge_id(u, v)
+        idx = self._graph.edge_id(u, v)
+        return idx.tousertensor() if force_multi or self.is_multigraph else idx[0]
 
-    def edge_ids(self, u, v):
+    def edge_ids(self, u, v, force_multi=False):
         """Return the edge ids.
 
         Parameters
@@ -273,16 +290,23 @@ def edge_ids(self, u, v):
             The src nodes.
         v : list, tensor
             The dst nodes.
+        force_multi : bool
+            If False, will return a single edge ID array if the graph is a simple graph.
+            If True, will always return 3 arrays (src nodes, dst nodes, edge ids).
 
         Returns
         -------
-        tensor
-            The edge id array.
+        tensor, or (tensor, tensor, tensor)
+        If force_multi is True or the graph is multigraph, return (src nodes, dst nodes, edge ids)
+        Otherwise, return a single tensor of edge ids.
         """
         u = utils.toindex(u)
         v = utils.toindex(v)
-        rst = self._graph.edge_ids(u, v)
-        return rst.tousertensor()
+        src, dst, eid = self._graph.edge_ids(u, v)
+        if force_multi or self.is_multigraph:
+            return src.tousertensor(), dst.tousertensor(), eid.tousertensor()
+        else:
+            return eid.tousertensor()
 
     def in_edges(self, v):
         """Return the in edges of the node(s).
@@ -612,7 +636,7 @@ def set_e_repr(self, h_uv, u=ALL, v=ALL):
         else:
             u = utils.toindex(u)
             v = utils.toindex(v)
-            eid = self._graph.edge_ids(u, v)
+            _, _, eid = self._graph.edge_ids(u, v)
             self.set_e_repr_by_id(h_uv, eid=eid)
 
     def set_e_repr_by_id(self, h_uv, eid=ALL):
@@ -674,7 +698,7 @@ def get_e_repr(self, u=ALL, v=ALL):
         else:
             u = utils.toindex(u)
             v = utils.toindex(v)
-            eid = self._graph.edge_ids(u, v)
+            _, _, eid = self._graph.edge_ids(u, v)
             return self.get_e_repr_by_id(eid=eid)
 
     def pop_e_repr(self, key=__REPR__):
@@ -809,28 +833,35 @@ def _apply_nodes(self, v, apply_node_func="default", reduce_accum=None):
             new_repr = reduce_accum
         self.set_n_repr(new_repr, v)
 
-    def apply_edges(self, u, v, apply_edge_func="default"):
+    def apply_edges(self, u=None, v=None, apply_edge_func="default", eid=None):
         """Apply the function on edge representations.
 
         Parameters
         ----------
-        u : int, iterable of int, tensor
+        u : optional, int, iterable of int, tensor
           The src node id(s).
-        v : int, iterable of int, tensor
+        v : optional, int, iterable of int, tensor
           The dst node id(s).
         apply_edge_func : callable
           The apply edge function.
+        eid : None, edge, container or tensor
+          The edge to update on.  If eid is not None then u and v are ignored.
         """
         if apply_edge_func == "default":
             apply_edge_func = self._apply_edge_func
         if not apply_edge_func:
             # Skip none function call.
             return
-        new_repr = apply_edge_func(self.get_e_repr(u, v))
-        self.set_e_repr(new_repr, u, v)
 
-    def send(self, u, v, message_func="default"):
-        """Trigger the message function on edge u->v
+        if eid is None:
+            new_repr = apply_edge_func(self.get_e_repr(u, v))
+            self.set_e_repr(new_repr, u, v)
+        else:
+            new_repr = apply_edge_func(self.get_e_repr_by_id(eid))
+            self.set_e_repr_by_id(new_repr, eid)
+
+    def send(self, u=None, v=None, message_func="default", eid=None):
+        """Trigger the message function on edge u->v or eid
 
         The message function should be compatible with following signature:
 
@@ -842,45 +873,105 @@ def send(self, u, v, message_func="default"):
         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.
+
         Parameters
         ----------
-        u : node, container or tensor
+        u : optional, node, container or tensor
           The source node(s).
-        v : node, container or tensor
+        v : optional, node, container or tensor
           The destination node(s).
         message_func : callable
           The message function.
+        eid : optional, edge, container or tensor
+          The edge to update on.  If eid is not None then u and v are ignored.
+
+        Notes
+        -----
+        On multigraphs, if u and v are specified, then the messages will be sent
+        along all edges between u and v.
         """
         if message_func == "default":
             message_func = self._message_func
         assert message_func is not None
         if isinstance(message_func, (tuple, list)):
             message_func = BundledMessageFunction(message_func)
-        self._batch_send(u, v, message_func)
+        self._batch_send(u, v, eid, message_func)
 
-    def _batch_send(self, u, v, message_func):
-        if is_all(u) and is_all(v):
-            u, v, _ = self._graph.edges()
-            self._msg_graph.add_edges(u, v) # TODO(minjie): can be optimized
+    def _batch_send(self, u, v, eid, message_func):
+        if is_all(u) and is_all(v) and eid is None:
+            u, v, eid = self._graph.edges()
             # call UDF
             src_reprs = self.get_n_repr(u)
             edge_reprs = self.get_e_repr()
             msgs = message_func(src_reprs, edge_reprs)
+        elif eid is not None:
+            eid = utils.toindex(eid)
+            u, v, _ = self._graph.find_edges(eid)
+            # call UDF
+            src_reprs = self.get_n_repr(u)
+            edge_reprs = self.get_e_repr_by_id(eid)
+            msgs = message_func(src_reprs, edge_reprs)
         else:
             u = utils.toindex(u)
             v = utils.toindex(v)
-            u, v = utils.edge_broadcasting(u, v)
-            self._msg_graph.add_edges(u, v)
+            u, v, eid = self._graph.edge_ids(u, v)
             # call UDF
             src_reprs = self.get_n_repr(u)
-            edge_reprs = self.get_e_repr(u, v)
+            edge_reprs = self.get_e_repr_by_id(eid)
             msgs = message_func(src_reprs, edge_reprs)
+
+        new_uv = []
+        msg_target_rows = []
+        msg_update_rows = []
+        msg_append_rows = []
+        for i, (_u, _v, _eid) in enumerate(zip(u, v, eid)):
+            if _eid in self._msg_edges:
+                msg_target_rows.append(self._msg_edges.index(_eid))
+                msg_update_rows.append(i)
+            else:
+                new_uv.append((_u, _v))
+                self._msg_edges.append(_eid)
+                msg_append_rows.append(i)
+
+        msg_target_rows = utils.toindex(msg_target_rows)
+        msg_update_rows = utils.toindex(msg_update_rows)
+        msg_append_rows = utils.toindex(msg_append_rows)
+
         if utils.is_dict_like(msgs):
-            self._msg_frame.append(msgs)
+            if len(msg_target_rows) > 0:
+                self._msg_frame.update_rows(
+                        msg_target_rows,
+                        {k: F.gather_row(msgs[k], msg_update_rows.tousertensor())
+                            for k in msgs}
+                        )
+            if len(msg_append_rows) > 0:
+                new_u, new_v = zip(*new_uv)
+                new_u = utils.toindex(new_u)
+                new_v = utils.toindex(new_v)
+                self._msg_graph.add_edges(new_u, new_v)
+                self._msg_frame.append(
+                        {k: F.gather_row(msgs[k], msg_append_rows.tousertensor())
+                            for k in msgs}
+                        )
         else:
-            self._msg_frame.append({__MSG__ : msgs})
-
-    def update_edge(self, u=ALL, v=ALL, edge_func="default"):
+            if len(msg_target_rows) > 0:
+                self._msg_frame.update_rows(
+                        msg_target_rows,
+                        {__MSG__: F.gather_row(msgs, msg_update_rows.tousertensor())}
+                        )
+            if len(msg_append_rows) > 0:
+                new_u, new_v = zip(*new_uv)
+                new_u = utils.toindex(new_u)
+                new_v = utils.toindex(new_v)
+                self._msg_graph.add_edges(new_u, new_v)
+                self._msg_frame.append(
+                        {__MSG__: F.gather_row(msgs, msg_append_rows.tousertensor())}
+                        )
+
+    def update_edge(self, u=ALL, v=ALL, edge_func="default", eid=None):
         """Update representation on edge u->v
 
         The edge function should be compatible with following signature:
@@ -899,15 +990,17 @@ def update_edge(self, u=ALL, v=ALL, edge_func="default"):
           The destination node(s).
         edge_func : callable
           The update function.
+        eid : optional, edge, container or tensor
+          The edge to update on.  If eid is not None then u and v are ignored.
         """
         if edge_func == "default":
             edge_func = self._edge_func
         assert edge_func is not None
-        self._batch_update_edge(u, v, edge_func)
+        self._batch_update_edge(u, v, eid, edge_func)
 
-    def _batch_update_edge(self, u, v, edge_func):
-        if is_all(u) and is_all(v):
-            u, v, _ = self._graph.edges()
+    def _batch_update_edge(self, u, v, eid, edge_func):
+        if is_all(u) and is_all(v) and eid is None:
+            u, v, eid = self._graph.edges()
             # call the UDF
             src_reprs = self.get_n_repr(u)
             dst_reprs = self.get_n_repr(v)
@@ -915,10 +1008,12 @@ def _batch_update_edge(self, u, v, edge_func):
             new_edge_reprs = edge_func(src_reprs, dst_reprs, edge_reprs)
             self.set_e_repr(new_edge_reprs)
         else:
-            u = utils.toindex(u)
-            v = utils.toindex(v)
-            u, v = utils.edge_broadcasting(u, v)
-            eid = self._graph.edge_ids(u, v)
+            if eid is None:
+                u = utils.toindex(u)
+                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)
@@ -1036,17 +1131,19 @@ def _reshape_fn(msg):
             self.set_n_repr(new_reprs, reordered_v)
 
     def send_and_recv(self,
-                      u, v,
+                      u=None, v=None,
                       message_func="default",
                       reduce_func="default",
-                      apply_node_func="default"):
-        """Trigger the message function on u->v and update v.
+                      apply_node_func="default",
+                      eid=None):
+        """Trigger the message function on u->v and update v, or on edge eid
+        and update the destination nodes.
 
         Parameters
         ----------
-        u : node, container or tensor
+        u : optional, node, container or tensor
           The source node(s).
-        v : node, container or tensor
+        v : optional, node, container or tensor
           The destination node(s).
         message_func : callable
           The message function.
@@ -1054,14 +1151,12 @@ def send_and_recv(self,
           The reduce function.
         apply_node_func : callable, optional
           The update function.
-        """
-        u = utils.toindex(u)
-        v = utils.toindex(v)
-        if len(u) == 0:
-            # no edges to be triggered
-            assert len(v) == 0
-            return
 
+        Notes
+        -----
+        On multigraphs, if u and v are specified, then the messages will be sent
+        and received along all edges between u and v.
+        """
         if message_func == "default":
             message_func = self._message_func
         if reduce_func == "default":
@@ -1069,14 +1164,42 @@ def send_and_recv(self,
         assert message_func is not None
         assert reduce_func is not None
 
-        executor = scheduler.get_executor(
-                'send_and_recv', self, src=u, dst=v,
-                message_func=message_func, reduce_func=reduce_func)
+        if eid is None:
+            if u is None or v is None:
+                raise ValueError('u and v must be given if eid is None')
+
+            u = utils.toindex(u)
+            v = utils.toindex(v)
+            if len(u) == 0:
+                # no edges to be triggered
+                assert len(v) == 0
+                return
+            unique_v = utils.toindex(F.unique(v.tousertensor()))
+
+            executor = scheduler.get_executor(
+                    'send_and_recv', self, src=u, dst=v,
+                    message_func=message_func, reduce_func=reduce_func)
+        else:
+            eid = utils.toindex(eid)
+            if len(eid) == 0:
+                # no edges to be triggered
+                return
+
+            executor = None
+
         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
+            self.send(eid=eid, message_func=message_func)
+            self.recv(unique_v, reduce_func, apply_node_func)
         else:
             # handle multiple message and reduce func
             if isinstance(message_func, (tuple, list)):
@@ -1103,7 +1226,7 @@ def send_and_recv(self,
             new_reprs = executor.run()
             unique_v = executor.recv_nodes
 
-        self._apply_nodes(unique_v, apply_node_func, reduce_accum=new_reprs)
+            self._apply_nodes(unique_v, apply_node_func, reduce_accum=new_reprs)
 
     def pull(self,
              v,
@@ -1238,8 +1361,25 @@ def subgraph(self, nodes):
             The subgraph.
         """
         induced_nodes = utils.toindex(nodes)
-        gi, induced_edges = self._graph.node_subgraph(induced_nodes)
-        return dgl.DGLSubGraph(self, induced_nodes, induced_edges, gi)
+        sgi = self._graph.node_subgraph(induced_nodes)
+        return dgl.DGLSubGraph(self, sgi.induced_nodes, sgi.induced_edges, sgi)
+
+    def edge_subgraph(self, edges):
+        """Generate the subgraph among the given edges.
+
+        Parameters
+        ----------
+        edges : list, or iterable
+            A container of the edges to construct subgraph.
+
+        Returns
+        -------
+        G : DGLSubGraph
+            The subgraph.
+        """
+        induced_edges = utils.toindex(edges)
+        sgi = self._graph.edge_subgraph(induced_edges)
+        return dgl.DGLSubGraph(self, sgi.induced_nodes, sgi.induced_edges, sgi)
 
     def merge(self, subgraphs, reduce_func='sum'):
         """Merge subgraph features back to this parent graph.
diff --git a/python/dgl/graph_index.py b/python/dgl/graph_index.py
--- a/python/dgl/graph_index.py
+++ b/python/dgl/graph_index.py
@@ -72,6 +72,16 @@ def clear(self):
         _CAPI_DGLGraphClear(self._handle)
         self._cache.clear()
 
+    def is_multigraph(self):
+        """Return whether the graph is a multigraph
+
+        Returns
+        -------
+        bool
+            True if it is a multigraph, False otherwise.
+        """
+        return bool(_CAPI_DGLGraphIsMultigraph(self._handle))
+
     def number_of_nodes(self):
         """Return the number of nodes.
 
@@ -103,9 +113,9 @@ def has_node(self, vid):
         Returns
         -------
         bool
-            True if the node exists
+            True if the node exists, False otherwise.
         """
-        return _CAPI_DGLGraphHasVertex(self._handle, vid)
+        return bool(_CAPI_DGLGraphHasVertex(self._handle, vid))
 
     def has_nodes(self, vids):
         """Return true if the nodes exist.
@@ -123,7 +133,7 @@ def has_nodes(self, vids):
         vid_array = vids.todgltensor()
         return utils.toindex(_CAPI_DGLGraphHasVertices(self._handle, vid_array))
 
-    def has_edge(self, u, v):
+    def has_edge_between(self, u, v):
         """Return true if the edge exists.
 
         Parameters
@@ -136,11 +146,11 @@ def has_edge(self, u, v):
         Returns
         -------
         bool
-            True if the edge exists
+            True if the edge exists, False otherwise
         """
-        return _CAPI_DGLGraphHasEdge(self._handle, u, v)
+        return bool(_CAPI_DGLGraphHasEdgeBetween(self._handle, u, v))
 
-    def has_edges(self, u, v):
+    def has_edges_between(self, u, v):
         """Return true if the edge exists.
 
         Parameters
@@ -157,7 +167,7 @@ def has_edges(self, u, v):
         """
         u_array = u.todgltensor()
         v_array = v.todgltensor()
-        return utils.toindex(_CAPI_DGLGraphHasEdges(self._handle, u_array, v_array))
+        return utils.toindex(_CAPI_DGLGraphHasEdgesBetween(self._handle, u_array, v_array))
 
     def predecessors(self, v, radius=1):
         """Return the predecessors of the node.
@@ -194,7 +204,7 @@ def successors(self, v, radius=1):
         return utils.toindex(_CAPI_DGLGraphSuccessors(self._handle, v, radius))
 
     def edge_id(self, u, v):
-        """Return the id of the edge.
+        """Return the id array of all edges between u and v.
 
         Parameters
         ----------
@@ -205,13 +215,13 @@ def edge_id(self, u, v):
 
         Returns
         -------
-        int
-            The edge id.
+        utils.Index
+            The edge id array.
         """
-        return _CAPI_DGLGraphEdgeId(self._handle, u, v)
+        return utils.toindex(_CAPI_DGLGraphEdgeId(self._handle, u, v))
 
     def edge_ids(self, u, v):
-        """Return the edge ids.
+        """Return a triplet of arrays that contains the edge IDs.
 
         Parameters
         ----------
@@ -223,11 +233,47 @@ def edge_ids(self, u, v):
         Returns
         -------
         utils.Index
-            Teh edge id array.
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
         """
         u_array = u.todgltensor()
         v_array = v.todgltensor()
-        return utils.toindex(_CAPI_DGLGraphEdgeIds(self._handle, u_array, v_array))
+        edge_array = _CAPI_DGLGraphEdgeIds(self._handle, u_array, v_array)
+
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
+
+        return src, dst, eid
+
+    def find_edges(self, eid):
+        """Return a triplet of arrays that contains the edge IDs.
+
+        Parameters
+        ----------
+        eid : utils.Index
+            The edge ids.
+
+        Returns
+        -------
+        utils.Index
+            The src nodes.
+        utils.Index
+            The dst nodes.
+        utils.Index
+            The edge ids.
+        """
+        eid_array = eid.todgltensor()
+        edge_array = _CAPI_DGLGraphFindEdges(self._handle, eid_array)
+
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
+
+        return src, dst, eid
 
     def in_edges(self, v):
         """Return the in edges of the node(s).
@@ -378,16 +424,32 @@ def node_subgraph(self, v):
 
         Returns
         -------
-        GraphIndex
+        SubgraphIndex
             The subgraph index.
-        utils.Index
-            The induced edge ids. This is also a map from new edge id to parent edge id.
         """
         v_array = v.todgltensor()
         rst = _CAPI_DGLGraphVertexSubgraph(self._handle, v_array)
-        gi = GraphIndex(rst(0))
         induced_edges = utils.toindex(rst(2))
-        return gi, induced_edges
+        return SubgraphIndex(rst(0), self, v, induced_edges)
+
+    def edge_subgraph(self, e):
+        """Return the induced edge subgraph.
+
+        Parameters
+        ----------
+        e : utils.Index
+            The edges.
+
+        Returns
+        -------
+        SubgraphIndex
+            The subgraph index.
+        """
+        e_array = e.todgltensor()
+        rst = _CAPI_DGLGraphEdgeSubgraph(self._handle, e_array)
+        gi = GraphIndex(rst(0))
+        induced_nodes = utils.toindex(rst(1))
+        return SubgraphIndex(rst(0), self, induced_nodes, e)
 
     def adjacency_matrix(self):
         """Return the adjacency matrix representation of this graph.
@@ -460,7 +522,7 @@ def to_networkx(self):
             The nx graph
         """
         src, dst, eid = self.edges()
-        ret = nx.DiGraph()
+        ret = nx.MultiDiGraph() if self.is_multigraph() else nx.DiGraph()
         for u, v, id in zip(src, dst, eid):
             ret.add_edge(u, v, id=id)
         return ret
@@ -477,8 +539,13 @@ def from_networkx(self, nx_graph):
             The nx graph
         """
         self.clear()
-        if not isinstance(nx_graph, nx.DiGraph):
-            nx_graph = nx.DiGraph(nx_graph)
+
+        if not isinstance(nx_graph, nx.Graph):
+            nx_graph = (nx.MultiDiGraph(nx_graph) if self.is_multigraph()
+                    else nx.DiGraph(nx_graph))
+        else:
+            nx_graph = nx_graph.to_directed()
+
         num_nodes = nx_graph.number_of_nodes()
         self.add_nodes(num_nodes)
         has_edge_id = 'id' in next(iter(nx_graph.edges))
@@ -487,16 +554,16 @@ def from_networkx(self, nx_graph):
             src = np.zeros((num_edges,), dtype=np.int64)
             dst = np.zeros((num_edges,), dtype=np.int64)
             for e, attr in nx_graph.edges.items:
-                u, v = e
+                # MultiDiGraph returns a triplet in e while DiGraph returns a pair
                 eid = attr['id']
-                src[eid] = u
-                dst[eid] = v
+                src[eid] = e[0]
+                dst[eid] = e[1]
         else:
             src = []
             dst = []
-            for u, v in nx_graph.edges:
-                src.append(u)
-                dst.append(v)
+            for e in nx_graph.edges:
+                src.append(e[0])
+                dst.append(e[1])
         src = utils.toindex(src)
         dst = utils.toindex(dst)
         self.add_edges(src, dst)
@@ -531,7 +598,32 @@ def line_graph(self, backtracking=True):
         """
         handle = _CAPI_DGLGraphLineGraph(self._handle, backtracking)
         return GraphIndex(handle)
-        
+
+class SubgraphIndex(GraphIndex):
+    def __init__(self, handle, parent, induced_nodes, induced_edges):
+        super().__init__(handle)
+
+        self._parent = parent
+        self._induced_nodes = induced_nodes
+        self._induced_edges = induced_edges
+
+    def add_nodes(self, num):
+        raise RuntimeError('Readonly graph. Mutation is not allowed.')
+
+    def add_edge(self, u, v):
+        raise RuntimeError('Readonly graph. Mutation is not allowed.')
+
+    def add_edges(self, u, v):
+        raise RuntimeError('Readonly graph. Mutation is not allowed.')
+
+    @property
+    def induced_edges(self):
+        return self._induced_edges
+
+    @property
+    def induced_nodes(self):
+        return self._induced_nodes
+
 def disjoint_union(graphs):
     """Return a disjoint union of the input graphs.
 
@@ -590,17 +682,20 @@ def disjoint_partition(graph, num_or_size_splits):
         graphs.append(GraphIndex(handle))
     return graphs
 
-def create_graph_index(graph_data=None):
+def create_graph_index(graph_data=None, multigraph=False):
     """Create a graph index object.
 
     Parameters
     ----------
     graph_data : graph data, optional
         Data to initialize graph. Same as networkx's semantics.
+    multigraph : bool, optional
+        Whether the graph is multigraph (default is False)
     """
     if isinstance(graph_data, GraphIndex):
         return graph_data
-    handle = _CAPI_DGLGraphCreate()
+
+    handle = _CAPI_DGLGraphCreate(multigraph)
     gi = GraphIndex(handle)
     if graph_data is not None:
         gi.from_networkx(graph_data)
diff --git a/python/dgl/scheduler.py b/python/dgl/scheduler.py
--- a/python/dgl/scheduler.py
+++ b/python/dgl/scheduler.py
@@ -416,6 +416,14 @@ def __init__(self, graph, src, dst, mfunc, rfunc):
         BundledExecutor.__init__(self, graph, mfunc, rfunc)
 
 def _is_spmv_supported(fn, graph=None):
+    # FIXME: also take into account
+    # (1) which backend DGL is under.
+    # (2) whether the graph is a multigraph.
+    #
+    # Current SPMV optimizer assumes that duplicate entries are summed up
+    # in sparse matrices, which is the case for PyTorch but not MXNet.
+    # The result is that on multigraphs, SPMV can still work for reducer=sum
+    # and message=copy_src/src_mul_edge *only in PyTorch*.
     if isinstance(fn, fmsg.MessageFunction):
         return fn.is_spmv_supported(graph)
     elif isinstance(fn, fred.ReduceFunction):


Test patch:
diff --git a/tests/graph_index/test_basics.py b/tests/graph_index/test_basics.py
new file mode 100644
--- /dev/null
+++ b/tests/graph_index/test_basics.py
@@ -0,0 +1,125 @@
+from dgl import DGLError
+from dgl.utils import toindex
+from dgl.graph_index import create_graph_index
+import networkx as nx
+
+def test_edge_id():
+    gi = create_graph_index(multigraph=False)
+    assert not gi.is_multigraph()
+
+    gi = create_graph_index(multigraph=True)
+
+    gi.add_nodes(4)
+    gi.add_edge(0, 1)
+    eid = gi.edge_id(0, 1).tolist()
+    assert len(eid) == 1
+    assert eid[0] == 0
+    assert gi.is_multigraph()
+
+    # multiedges
+    gi.add_edge(0, 1)
+    eid = gi.edge_id(0, 1).tolist()
+    assert len(eid) == 2
+    assert eid[0] == 0
+    assert eid[1] == 1
+
+    gi.add_edges(toindex([0, 1, 1, 2]), toindex([2, 2, 2, 3]))
+    src, dst, eid = gi.edge_ids(toindex([0, 0, 2, 1]), toindex([2, 1, 3, 2]))
+    eid_answer = [2, 0, 1, 5, 3, 4]
+    assert len(eid) == 6
+    assert all(e == ea for e, ea in zip(eid, eid_answer))
+
+    # find edges
+    src, dst, eid = gi.find_edges(toindex([1, 3, 5]))
+    assert len(src) == len(dst) == len(eid) == 3
+    assert src[0] == 0 and src[1] == 1 and src[2] == 2
+    assert dst[0] == 1 and dst[1] == 2 and dst[2] == 3
+    assert eid[0] == 1 and eid[1] == 3 and eid[2] == 5
+
+    # source broadcasting
+    src, dst, eid = gi.edge_ids(toindex([0]), toindex([1, 2]))
+    eid_answer = [0, 1, 2]
+    assert len(eid) == 3
+    assert all(e == ea for e, ea in zip(eid, eid_answer))
+
+    # destination broadcasting
+    src, dst, eid = gi.edge_ids(toindex([1, 0]), toindex([2]))
+    eid_answer = [3, 4, 2]
+    assert len(eid) == 3
+    assert all(e == ea for e, ea in zip(eid, eid_answer))
+
+    gi.clear()
+    # the following assumes that grabbing nonexistent edge will throw an error
+    try:
+        gi.edge_id(0, 1)
+        fail = True
+    except DGLError:
+        fail = False
+    finally:
+        assert not fail
+
+    gi.add_nodes(4)
+    gi.add_edge(0, 1)
+    eid = gi.edge_id(0, 1).tolist()
+    assert len(eid) == 1
+    assert eid[0] == 0
+
+def test_nx():
+    gi = create_graph_index(multigraph=True)
+
+    gi.add_nodes(2)
+    gi.add_edge(0, 1)
+    nxg = gi.to_networkx()
+    assert len(nxg.nodes) == 2
+    assert len(nxg.edges(0, 1)) == 1
+    gi.add_edge(0, 1)
+    nxg = gi.to_networkx()
+    assert len(nxg.edges(0, 1)) == 2
+
+    nxg = nx.DiGraph()
+    nxg.add_edge(0, 1)
+    gi = create_graph_index(nxg)
+    assert not gi.is_multigraph()
+    assert gi.number_of_nodes() == 2
+    assert gi.number_of_edges() == 1
+    assert gi.edge_id(0, 1)[0] == 0
+
+    nxg = nx.MultiDiGraph()
+    nxg.add_edge(0, 1)
+    nxg.add_edge(0, 1)
+    gi = create_graph_index(nxg, True)
+    assert gi.is_multigraph()
+    assert gi.number_of_nodes() == 2
+    assert gi.number_of_edges() == 2
+    assert 0 in gi.edge_id(0, 1)
+    assert 1 in gi.edge_id(0, 1)
+
+def test_predsucc():
+    gi = create_graph_index(multigraph=True)
+
+    gi.add_nodes(4)
+    gi.add_edge(0, 1)
+    gi.add_edge(0, 1)
+    gi.add_edge(0, 2)
+    gi.add_edge(2, 0)
+    gi.add_edge(3, 0)
+    gi.add_edge(0, 0)
+    gi.add_edge(0, 0)
+
+    pred = gi.predecessors(0)
+    assert len(pred) == 3
+    assert 2 in pred
+    assert 3 in pred
+    assert 0 in pred
+
+    succ = gi.successors(0)
+    assert len(succ) == 3
+    assert 1 in succ
+    assert 2 in succ
+    assert 0 in succ
+
+
+if __name__ == '__main__':
+    test_edge_id()
+    test_nx()
+    test_predsucc()
diff --git a/tests/graph_index/test_subgraph.py b/tests/graph_index/test_subgraph.py
new file mode 100644
--- /dev/null
+++ b/tests/graph_index/test_subgraph.py
@@ -0,0 +1,38 @@
+from dgl import DGLError
+from dgl.utils import toindex
+from dgl.graph_index import create_graph_index
+
+def test_node_subgraph():
+    gi = create_graph_index()
+    gi.add_nodes(4)
+    gi.add_edge(0, 1)
+    gi.add_edge(0, 2)
+    gi.add_edge(0, 2)
+    gi.add_edge(0, 3)
+
+    sub2par_nodemap = [2, 0, 3]
+    sgi = gi.node_subgraph(toindex(sub2par_nodemap))
+
+    for s, d, e in zip(*sgi.edges()):
+        assert sgi.induced_edges[e] in gi.edge_id(
+                sgi.induced_nodes[s], sgi.induced_nodes[d])
+
+def test_edge_subgraph():
+    gi = create_graph_index()
+    gi.add_nodes(4)
+    gi.add_edge(0, 1)
+    gi.add_edge(0, 1)
+    gi.add_edge(0, 2)
+    gi.add_edge(2, 3)
+
+    sub2par_edgemap = [3, 2]
+    sgi = gi.edge_subgraph(toindex(sub2par_edgemap))
+
+    for s, d, e in zip(*sgi.edges()):
+        assert sgi.induced_edges[e] in gi.edge_id(
+                sgi.induced_nodes[s], sgi.induced_nodes[d])
+
+
+if __name__ == '__main__':
+    test_node_subgraph()
+    test_edge_subgraph()
diff --git a/tests/pytorch/test_basics.py b/tests/pytorch/test_basics.py
--- a/tests/pytorch/test_basics.py
+++ b/tests/pytorch/test_basics.py
@@ -242,6 +242,104 @@ def _reduce(node, msgs):
     assert th.allclose(new_repr[0], old_repr[0])
     assert th.allclose(new_repr[1], old_repr[0])
 
+def test_send_twice():
+    g = DGLGraph()
+    g.add_nodes(3)
+    g.add_edge(0, 1)
+    g.add_edge(2, 1)
+    def _message_a(src, edge):
+        return {'a': src['a']}
+    def _message_b(src, edge):
+        return {'a': src['a'] * 3}
+    def _reduce(node, msgs):
+        assert msgs is not None
+        return {'a': msgs['a'].max(1)[0]}
+
+    old_repr = th.randn(3, 5)
+    g.set_n_repr({'a': old_repr})
+    g.send(0, 1, _message_a)
+    g.send(0, 1, _message_b)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], old_repr[0] * 3)
+
+    g.set_n_repr({'a': old_repr})
+    g.send(0, 1, _message_a)
+    g.send(2, 1, _message_b)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], th.stack([old_repr[0], old_repr[2] * 3], 0).max(0)[0])
+
+def test_send_multigraph():
+    g = DGLGraph(multigraph=True)
+    g.add_nodes(3)
+    g.add_edge(0, 1)
+    g.add_edge(0, 1)
+    g.add_edge(0, 1)
+    g.add_edge(2, 1)
+
+    def _message_a(src, edge):
+        return {'a': edge['a']}
+    def _message_b(src, edge):
+        return {'a': edge['a'] * 3}
+    def _reduce(node, msgs):
+        assert msgs is not None
+        return {'a': msgs['a'].max(1)[0]}
+
+    def answer(*args):
+        return th.stack(args, 0).max(0)[0]
+
+    # send by eid
+    old_repr = th.randn(4, 5)
+    g.set_n_repr({'a': th.zeros(3, 5)})
+    g.set_e_repr({'a': old_repr})
+    g.send(eid=[0, 2], message_func=_message_a)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], answer(old_repr[0], old_repr[2]))
+
+    g.set_n_repr({'a': th.zeros(3, 5)})
+    g.set_e_repr({'a': old_repr})
+    g.send(eid=[0, 2, 3], message_func=_message_a)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
+
+    # send on multigraph
+    g.set_n_repr({'a': th.zeros(3, 5)})
+    g.set_e_repr({'a': old_repr})
+    g.send([0, 2], [1, 1], _message_a)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], old_repr.max(0)[0])
+
+    # consecutive send and send_on
+    g.set_n_repr({'a': th.zeros(3, 5)})
+    g.set_e_repr({'a': old_repr})
+    g.send(2, 1, _message_a)
+    g.send(eid=[0, 1], message_func=_message_b)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], answer(old_repr[0] * 3, old_repr[1] * 3, old_repr[3]))
+
+    # consecutive send_on
+    g.set_n_repr({'a': th.zeros(3, 5)})
+    g.set_e_repr({'a': old_repr})
+    g.send(eid=0, message_func=_message_a)
+    g.send(eid=1, message_func=_message_b)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], answer(old_repr[0], old_repr[1] * 3))
+
+    # send_and_recv_on
+    g.set_n_repr({'a': th.zeros(3, 5)})
+    g.set_e_repr({'a': old_repr})
+    g.send_and_recv(eid=[0, 2, 3], message_func=_message_a, reduce_func=_reduce)
+    new_repr = g.get_n_repr()['a']
+    assert th.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
+    assert th.allclose(new_repr[[0, 2]], th.zeros(2, 5))
+
+
 if __name__ == '__main__':
     test_batch_setter_getter()
     test_batch_setter_autograd()
@@ -250,3 +348,5 @@ def _reduce(node, msgs):
     test_update_routines()
     test_reduce_0deg()
     test_pull_0deg()
+    test_send_twice()
+    test_send_multigraph()
diff --git a/tests/pytorch/test_line_graph.py b/tests/pytorch/test_line_graph.py
--- a/tests/pytorch/test_line_graph.py
+++ b/tests/pytorch/test_line_graph.py
@@ -39,8 +39,8 @@ def test_no_backtracking():
     for i in range(1, N):
         e1 = G.edge_id(0, i)
         e2 = G.edge_id(i, 0)
-        assert not L.has_edge(e1, e2)
-        assert not L.has_edge(e2, e1)
+        assert not L.has_edge_between(e1, e2)
+        assert not L.has_edge_between(e2, e1)
 
 if __name__ == '__main__':
     test_line_graph()