Repository: dmlc/dgl

Issue:
Messages not ordered by edge IDs in degree bucketing
## 🐛 Bug

See the following code.

## To Reproduce

```
In [1]: import dgl
Using backend: pytorch
INFO:rdflib:RDFLib Version: 4.2.2

In [3]: g = dgl.graph(([1, 3, 5, 0, 4, 2, 3, 3, 4, 5], [1, 1, 0, 0, 1, 2, 2, 0, 3, 3]))

In [5]: def reducer(nodes):
   ...:     print(nodes.mailbox['eid'])
   ...:     return {'n': nodes.mailbox['eid'].sum(1)}
   ...:

In [6]: import torch

In [7]: g.edata['eid'] = torch.arange(10)

In [8]: g.update_all(dgl.function.copy_e('eid', 'eid'), reducer)
tensor([[5, 6],
        [8, 9]])
tensor([[3, 7, 2],
        [0, 1, 4]])
```

In 0.4.3post2:

```
In [1]: import dgl
Using backend: pytorch

In [2]: g = dgl.graph(([1, 3, 5, 0, 4, 2, 3, 3, 4, 5], [1, 1, 0, 0, 1, 2, 2, 0, 3, 3]))

In [3]: def reducer(nodes):
   ...:     print(nodes.mailbox['eid'])
   ...:     return {'n': nodes.mailbox['eid'].sum(1)}

In [8]: g.edata['eid'] = torch.arange(10).float()

In [9]: g.update_all(dgl.function.copy_e('eid', 'eid'), reducer)
tensor([[2., 3., 7.],
        [0., 1., 4.]])
tensor([[8., 9.],
        [5., 6.]])
```

This inconsistency will break TreeLSTM.

## Expected behavior

Should be consistent as 0.4.3.

## Environment

 - DGL Version (e.g., 1.0): 0.5.1
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
 - OS (e.g., Linux):
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:



Discussion:

Resolving patch:
diff --git a/python/dgl/core.py b/python/dgl/core.py
--- a/python/dgl/core.py
+++ b/python/dgl/core.py
@@ -1,5 +1,6 @@
 """Implementation for core graph computation."""
 # pylint: disable=not-callable
+import numpy as np
 
 from .base import DGLError, is_all, NID, EID, ALL
 from . import backend as F
@@ -121,8 +122,13 @@ def invoke_udf_reduce(graph, func, msgdata, *, orig_nid=None):
             continue
         bkt_nodes.append(node_bkt)
         ndata_bkt = dstdata.subframe(node_bkt)
-        eid_bkt = graph.in_edges(node_bkt, form='eid')
+
+        # order the incoming edges per node by edge ID
+        eid_bkt = F.zerocopy_to_numpy(graph.in_edges(node_bkt, form='eid'))
         assert len(eid_bkt) == deg * len(node_bkt)
+        eid_bkt = np.sort(eid_bkt.reshape((len(node_bkt), deg)), 1)
+        eid_bkt = F.zerocopy_from_numpy(eid_bkt.flatten())
+
         msgdata_bkt = msgdata.subframe(eid_bkt)
         # reshape all msg tensors to (num_nodes_bkt, degree, feat_size)
         maildata = {}


Test patch:
diff --git a/tests/compute/test_basics.py b/tests/compute/test_basics.py
--- a/tests/compute/test_basics.py
+++ b/tests/compute/test_basics.py
@@ -644,3 +644,16 @@ def test_issue_1088(idtype):
     g = dgl.heterograph({('U', 'E', 'V'): ([0, 1, 2], [1, 2, 3])}, idtype=idtype, device=F.ctx())
     g.nodes['U'].data['x'] = F.randn((3, 3))
     g.update_all(fn.copy_u('x', 'm'), fn.sum('m', 'y'))
+
+@parametrize_dtype
+def test_degree_bucket_edge_ordering(idtype):
+    import dgl.function as fn
+    g = dgl.graph(
+        ([1, 3, 5, 0, 4, 2, 3, 3, 4, 5], [1, 1, 0, 0, 1, 2, 2, 0, 3, 3]),
+        idtype=idtype, device=F.ctx())
+    g.edata['eid'] = F.copy_to(F.arange(0, 10), F.ctx())
+    def reducer(nodes):
+        eid = F.asnumpy(F.copy_to(nodes.mailbox['eid'], F.cpu()))
+        assert np.array_equal(eid, np.sort(eid, 1))
+        return {'n': F.sum(nodes.mailbox['eid'], 1)}
+    g.update_all(fn.copy_e('eid', 'eid'), reducer)