[Bug] Fix #1088 (#1089)

* [Bug] Fix #1088

* fix

* add comment

src/graph/heterograph.cc

@@ -468,12 +468,13 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetRelationGraph")
 .set_body([] (DGLArgs args, DGLRetValue* rv) {
     HeteroGraphRef hg = args[0];
     dgl_type_t etype = args[1];
-    if (hg->NumEdgeTypes() == 1) {
-      CHECK_EQ(etype, 0);
-      *rv = hg;
-    } else {
+    CHECK_LE(etype, hg->NumEdgeTypes()) << "invalid edge type " << etype;
+    // Test if the heterograph is a unit graph.  If so, return itself.
+    auto bg = std::dynamic_pointer_cast<UnitGraph>(hg.sptr());
+    if (bg != nullptr)
+      *rv = bg;
+    else
       *rv = HeteroGraphRef(hg->GetRelationGraph(etype));
-    }
   });
 
 DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFlattenedGraph")

tests/compute/test_hetero_basics.py

@@ -762,6 +762,14 @@ def test_local_scope():
             assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
     foo(g)
 
+def test_issue_1088():
+    # This test ensures that message passing on a heterograph with one edge type
+    # would not crash (GitHub issue #1088).
+    import dgl.function as fn
+    g = dgl.heterograph({('U', 'E', 'V'): ([0, 1, 2], [1, 2, 3])})
+    g.nodes['U'].data['x'] = F.randn((3, 3))
+    g.update_all(fn.copy_u('x', 'm'), fn.sum('m', 'y'))
+
 if __name__ == '__main__':
     test_nx_conversion()
     test_batch_setter_getter()
@@ -781,3 +789,4 @@ if __name__ == '__main__':
     test_group_apply_edges()
     test_local_var()
     test_local_scope()
+    test_issue_1088()