[bugfix] Fix several issues in `test_kernel` (#801)

* upd

* stable

tests/compute/test_kernel.py

@@ -20,7 +20,6 @@ def udf_mean(nodes):
 def udf_sum(nodes):
     return {'r2': nodes.mailbox['m'].sum(1)}
 
-
 def udf_max(nodes):
     return {'r2': F.max(nodes.mailbox['m'], 1)}
 
@@ -40,27 +39,30 @@ def generate_feature(g, broadcast='none'):
     nv = g.number_of_nodes()
     ne = g.number_of_edges()
     if broadcast == 'e':
-        u = F.tensor(np.random.randn(nv, D1, D2, D3) + 1)
-        e = F.tensor(np.random.randn(ne, D2, 1) - 1)
-        v = F.tensor(np.random.randn(nv, D1, D2, D3))
+        u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+        e = F.tensor(np.random.uniform(-1, 1, (ne, D2, 1)))
+        v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
     elif broadcast == 'u':
-        u = F.tensor(np.random.randn(nv, D2, 1) + 1)
-        e = F.tensor(np.random.randn(ne, D1, D2, D3) - 1)
-        v = F.tensor(np.random.randn(nv, D1, D2, D3))
+        u = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1)))
+        e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
+        v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
     elif broadcast == 'v':
-        u = F.tensor(np.random.randn(nv, D1, D2, D3) + 1)
-        e = F.tensor(np.random.randn(ne, D1, D2, D3) - 1)
-        v = F.tensor(np.random.randn(nv, D2, 1))
+        u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+        e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
+        v = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1)))
     else:
-        u = F.tensor(np.random.randn(nv, D1, D2, D3) + 1)
-        e = F.tensor(np.random.randn(ne, D1, D2, D3) - 1)
-        v = F.tensor(np.random.randn(nv, D1, D2, D3))
+        u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+        e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
+        v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
     return u, v, e
 
 
 def test_copy_src_reduce():
     def _test(red, partial):
         g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
+        # NOTE(zihao): add self-loop to avoid zero-degree nodes.
+        # https://github.com/dmlc/dgl/issues/761
+        g.add_edges(g.nodes(), g.nodes())
         hu, hv, he = generate_feature(g, 'none')
         if partial:
             nid = F.tensor(list(range(0, 100, 2)))
@@ -110,6 +112,8 @@ def test_copy_src_reduce():
 def test_copy_edge_reduce():
     def _test(red, partial):
         g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
+        # NOTE(zihao): add self-loop to avoid zero-degree nodes.
+        g.add_edges(g.nodes(), g.nodes())
         hu, hv, he = generate_feature(g, 'none')
         if partial:
             nid = F.tensor(list(range(0, 100, 2)))
@@ -155,8 +159,30 @@ def test_copy_edge_reduce():
 
 
 def test_all_binary_builtins():
-    def _test(g, lhs, rhs, binary_op, reducer, paritial, nid, broadcast='none'):
+    def _test(g, lhs, rhs, binary_op, reducer, partial, nid, broadcast='none'):
+        # initialize node/edge features with uniform(-1, 1)
         hu, hv, he = generate_feature(g, broadcast)
+        if binary_op == 'div':
+            # op = div
+            # lhs range: [-1, 1]
+            # rhs range: [1, 2]
+            # result range: [-1, 1]
+            if rhs == 'u':
+                hu = (hu + 3) / 2
+            elif rhs == 'v':
+                hv = (hv + 3) / 2
+            elif rhs == 'e':
+                he = (he + 3) / 2
+
+        if binary_op == 'add' or binary_op == 'sub':
+            # op = add, sub
+            # lhs range: [-1/2, 1/2]
+            # rhs range: [-1/2, 1/2]
+            # result range: [-1, 1]
+            hu = hu / 2
+            hv = hv / 2
+            he = he / 2
+
         g.ndata['u'] = F.attach_grad(F.clone(hu))
         g.ndata['v'] = F.attach_grad(F.clone(hv))
         g.edata['e'] = F.attach_grad(F.clone(he))
@@ -200,9 +226,15 @@ def test_all_binary_builtins():
 
         def mfunc(edges):
             op = getattr(F, binary_op)
-            lhs_data = target_switch(edges, lhs)
-            rhs_data = target_switch(edges, rhs)
-            return {"m": op(lhs_data[lhs], rhs_data[rhs])}
+            lhs_data = target_switch(edges, lhs)[lhs]
+            rhs_data = target_switch(edges, rhs)[rhs]
+            # NOTE(zihao): we need to do batched broadcast
+            # e.g. (68, 3, 1) op (68, 5, 3, 4)
+            while F.ndim(lhs_data) < F.ndim(rhs_data):
+                lhs_data = F.unsqueeze(lhs_data, 1)
+            while F.ndim(rhs_data) < F.ndim(lhs_data):
+                rhs_data = F.unsqueeze(rhs_data, 1)
+            return {"m": op(lhs_data, rhs_data)}
 
         def rfunc(nodes):
             op = getattr(F, reducer)
@@ -249,6 +281,8 @@ def test_all_binary_builtins():
 
     g = dgl.DGLGraph()
     g.add_nodes(20)
+    # NOTE(zihao): add self-loop to avoid zero-degree nodes.
+    g.add_edges(g.nodes(), g.nodes())
     for i in range(2, 18):
         g.add_edge(0, i)
         g.add_edge(1, i)
@@ -267,7 +301,8 @@ def test_all_binary_builtins():
             for reducer in ["sum", "max", "min", "prod", "mean"]:
                 for broadcast in ["none", lhs, rhs]:
                     for partial in [False, True]:
-                        _test(g, lhs, rhs, binary_op, reducer, partial, nid)
+                        _test(g, lhs, rhs, binary_op, reducer, partial, nid,
+                              broadcast=broadcast)
 
 if __name__ == '__main__':
     test_copy_src_reduce()