[BugFix] Fix #1779 (#1781)

* Fix bug #1779

* Fix
Co-authored-by: Ubuntu <ubuntu@ip-172-31-51-214.ec2.internal>

python/dgl/batched_heterograph.py

@@ -284,11 +284,11 @@ class BatchedDGLHeteroGraph(DGLHeteroGraph):
         new_nframes = []
         for nframe in self._node_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in nframe.items()}
-            new_nframes.append(FrameRef(Frame(new_feats)))
+            new_nframes.append(FrameRef(Frame(new_feats, num_rows=nframe.num_rows)))
         new_eframes = []
         for eframe in self._edge_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in eframe.items()}
-            new_eframes.append(FrameRef(Frame(new_feats)))
+            new_eframes.append(FrameRef(Frame(new_feats, num_rows=eframe.num_rows)))
         # TODO(minjie): replace the following line with the commented one to enable GPU graph.
         new_gidx = self._graph
         #new_gidx = self._graph.copy_to(utils.to_dgl_context(ctx))

python/dgl/data/heterograph_serialize.py

@@ -49,12 +49,12 @@ class HeteroGraphData(ObjectBase):
         gidx = _CAPI_GetGindexFromHeteroGraphData(self)
         nframes = []
         eframes = []
-        for ntensor in ntensor_list:
+        for ntid, ntensor in enumerate(ntensor_list):
             ndict = {ntensor[i]: F.zerocopy_from_dgl_ndarray(ntensor[i+1]) for i in range(0, len(ntensor), 2)}
-            nframes.append(FrameRef(Frame(ndict)))
+            nframes.append(FrameRef(Frame(ndict, num_rows=gidx.number_of_nodes(ntid))))
         
-        for etensor in etensor_list:
+        for etid, etensor in enumerate(etensor_list):
             edict = {etensor[i]: F.zerocopy_from_dgl_ndarray(etensor[i+1]) for i in range(0, len(etensor), 2)}
-            eframes.append(FrameRef(Frame(edict)))
+            eframes.append(FrameRef(Frame(edict, num_rows=gidx.number_of_edges(etid))))
         
         return DGLHeteroGraph(gidx, ntype_names, etype_names, nframes, eframes)

python/dgl/heterograph.py

@@ -4316,11 +4316,11 @@ class DGLHeteroGraph(object):
         new_nframes = []
         for nframe in self._node_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in nframe.items()}
-            new_nframes.append(FrameRef(Frame(new_feats)))
+            new_nframes.append(FrameRef(Frame(new_feats, num_rows=nframe.num_rows)))
         new_eframes = []
         for eframe in self._edge_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in eframe.items()}
-            new_eframes.append(FrameRef(Frame(new_feats)))
+            new_eframes.append(FrameRef(Frame(new_feats, num_rows=eframe.num_rows)))
         # TODO(minjie): replace the following line with the commented one to enable GPU graph.
         new_gidx = self._graph
         #new_gidx = self._graph.copy_to(utils.to_dgl_context(ctx))

tests/compute/test_batched_heterograph.py

@@ -402,9 +402,22 @@ def test_to_device(index_dtype):
         assert bg.batch_num_nodes('user') == bg1.batch_num_nodes('user')
         assert bg.batch_num_edges('plays') == bg1.batch_num_edges('plays')
 
+    # set feature
+    g1 = dgl.heterograph({
+        ('user', 'plays', 'game'): [(0, 0), (1, 1)]
+    }, index_dtype=index_dtype)
+    g2 = dgl.heterograph({
+        ('user', 'plays', 'game'): [(0, 0), (1, 0)]
+    }, index_dtype=index_dtype)
+    bg = dgl.batch_hetero([g1, g2])
+    if F.is_cuda_available():
+        bg1 = bg.to(F.cuda())
+        bg1.nodes['user'].data['test'] = F.copy_to(F.tensor([0,1,2,3]), F.cuda())
+        bg1.edata['test'] = F.copy_to(F.tensor([0,1,2,3]), F.cuda())
+
 if __name__ == '__main__':
     test_batching_hetero_topology('int32')
     test_batching_hetero_and_batched_hetero_topology('int32')
     test_batched_features('int32')
     test_batching_with_zero_nodes_edges('int32')
-    # test_to_device()
+    test_to_device('int32')

tests/compute/test_heterograph.py

@@ -885,6 +885,15 @@ def test_to_device(index_dtype):
         g1 = g.to(F.cuda())
         assert g1 is not None
 
+    # set feature after g.to
+    g = create_test_heterograph(index_dtype)
+    if F.is_cuda_available():
+        g1 = g.to(F.cuda())
+        assert g1 is not None
+        g1.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cuda())
+        g1.nodes['game'].data['i'] = F.copy_to(F.ones((2, 5)), F.cuda())
+        g1.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cuda())
+
 @parametrize_dtype
 def test_convert_bound(index_dtype):
     def _test_bipartite_bound(data, card):
@@ -2068,7 +2077,7 @@ if __name__ == '__main__':
     # test_flatten()
     # test_convert_bound()
     # test_convert()
-    # test_to_device()
+    # test_to_device("int32")
     # test_transform("int32")
     # test_subgraph("int32")
     # test_subgraph_mask("int32")

tests/compute/test_serialize.py

@@ -233,7 +233,7 @@ def create_heterographs2(index_dtype):
     g_y.nodes['user'].data['hh'] = F.ones((4, 5))
     g_y.edges['knows'].data['ww'] = F.randn((2, 10))
     g = dgl.hetero_from_relations([g_x, g_y, g_z])
-    return [g, g_x, g_y]
+    return [g, g_x, g_y, g_z]
 
 def test_deserialize_old_heterograph_file():
     path = os.path.join(
@@ -273,11 +273,11 @@ def test_serialize_heterograph():
             assert g_list[i].canonical_etypes[j] == etypes
     assert g_list[1].restrict_format() == 'any'
     assert g_list[2].restrict_format() == 'csr'
-    assert g_list[3].idtype == F.int32
+    assert g_list[4].idtype == F.int32
     assert np.allclose(
         F.asnumpy(g_list[2].nodes['user'].data['hh']), np.ones((4, 5)))
     assert np.allclose(
-        F.asnumpy(g_list[5].nodes['user'].data['hh']), np.ones((4, 5)))
+        F.asnumpy(g_list[6].nodes['user'].data['hh']), np.ones((4, 5)))
     edges = g_list[0]['follows'].edges()
     assert np.allclose(F.asnumpy(edges[0]), np.array([0, 1, 2]))
     assert np.allclose(F.asnumpy(edges[1]), np.array([1, 2, 3]))
@@ -285,6 +285,10 @@ def test_serialize_heterograph():
         assert g_list[i].ntypes == g_list0[i].ntypes
         assert g_list[i].etypes == g_list0[i].etypes
 
+    # test set feature after load_graph
+    g_list[3].nodes['user'].data['test'] = F.tensor([0, 1, 2, 4])
+    g_list[3].edata['test'] = F.tensor([0, 1, 2])
+
     os.unlink(path)
 
 @pytest.mark.skip(reason="lack of permission on CI")
@@ -308,17 +312,17 @@ def test_serialize_heterograph_s3():
 
 if __name__ == "__main__":
     pass
-    test_graph_serialize_with_feature(True)
-    test_graph_serialize_with_feature(False)
-    test_graph_serialize_without_feature(True)
-    test_graph_serialize_without_feature(False)
-    test_graph_serialize_with_labels(True)
-    test_graph_serialize_with_labels(False)
-    test_serialize_tensors()
-    test_serialize_empty_dict()
-    test_load_old_files1()
-    test_load_old_files2()
-    test_serialize_heterograph()
-    # test_serialize_heterograph_s3()
-    test_deserialize_old_heterograph_file()
-    # create_old_heterograph_files()
+    #test_graph_serialize_with_feature(True)
+    #test_graph_serialize_with_feature(False)
+    #test_graph_serialize_without_feature(True)
+    #test_graph_serialize_without_feature(False)
+    #test_graph_serialize_with_labels(True)
+    #test_graph_serialize_with_labels(False)
+    #test_serialize_tensors()
+    #test_serialize_empty_dict()
+    #test_load_old_files1()
+    #test_load_old_files2()
+    #test_serialize_heterograph()
+    #test_serialize_heterograph_s3()
+    #test_deserialize_old_heterograph_file()
+    #create_old_heterograph_files()