[Doc] Update docs to motivate the use of set_batch_num_nodes and set_batch_num_edges (#2630)

* fix doc

* update docstring for transformation on batch graph

* update cn guide

* fix format

* update rst

docs/source/guide/training-graph.rst

@@ -80,6 +80,10 @@ about the list.
     bg.edges()
     # (tensor([0, 1, 2, 4, 4, 4, 5], tensor([1, 2, 3, 4, 5, 6, 4]))
 
+Please note that most dgl transformation functions will discard the batch information.
+In order to maintain such information, please use :func:`dgl.DGLGraph.set_batch_num_nodes`
+and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph.
+
 Graph Readout
 ^^^^^^^^^^^^^
 

docs/source/guide_cn/training-graph.rst

@@ -42,6 +42,9 @@
 
    批次化的图
 
+需要注意，DGL里对图进行变换的函数会去掉图上的批次信息。用户可以通过 :func:`dgl.DGLGraph.set_batch_num_nodes`
+和 :func:`dgl.DGLGraph.set_batch_num_edges` 两个函数在变换后的图上重新加入批次信息。
+
 图读出
 ^^^^^^^^^^^^^
 

python/dgl/heterograph.py

@@ -216,11 +216,15 @@ class DGLHeteroGraph(object):
 
         * Inplace update is applied to the current graph.
         * If the key of ``data`` does not contain some existing feature fields,
-        those features for the new nodes will be created by initializers
-        defined with :func:`set_n_initializer` (default initializer fills zeros).
+          those features for the new nodes will be created by initializers
+          defined with :func:`set_n_initializer` (default initializer fills zeros).
         * If the key of ``data`` contains new feature fields, those features for
-        the old nodes will be created by initializers defined with
-        :func:`set_n_initializer` (default initializer fills zeros).
+          the old nodes will be created by initializers defined with
+          :func:`set_n_initializer` (default initializer fills zeros).
+        * This function discards the batch information. Please use
+          :func:`dgl.DGLGraph.set_batch_num_nodes`
+          and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+          to maintain the information.
 
         Examples
         --------
@@ -366,16 +370,20 @@ class DGLHeteroGraph(object):
 
         * Inplace update is applied to the current graph.
         * If end nodes of adding edges does not exists, add_nodes is invoked
-        to add new nodes. The node features of the new nodes will be created
-        by initializers defined with :func:`set_n_initializer` (default
-        initializer fills zeros). In certain cases, it is recommanded to
-        add_nodes first and then add_edges.
+          to add new nodes. The node features of the new nodes will be created
+          by initializers defined with :func:`set_n_initializer` (default
+          initializer fills zeros). In certain cases, it is recommanded to
+          add_nodes first and then add_edges.
         * If the key of ``data`` does not contain some existing feature fields,
-        those features for the new edges will be created by initializers
-        defined with :func:`set_n_initializer` (default initializer fills zeros).
+          those features for the new edges will be created by initializers
+          defined with :func:`set_n_initializer` (default initializer fills zeros).
         * If the key of ``data`` contains new feature fields, those features for
-        the old edges will be created by initializers defined with
-        :func:`set_n_initializer` (default initializer fills zeros).
+          the old edges will be created by initializers defined with
+          :func:`set_n_initializer` (default initializer fills zeros).
+        * This function discards the batch information. Please use
+          :func:`dgl.DGLGraph.set_batch_num_nodes`
+          and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+          to maintain the information.
 
         Examples
         --------
@@ -541,6 +549,14 @@ class DGLHeteroGraph(object):
             and ``edata`` of the resulting graph under name ``dgl.NID`` and ``dgl.EID``,
             respectively.
 
+        Notes
+        -----
+
+        This function discards the batch information. Please use
+        :func:`dgl.DGLGraph.set_batch_num_nodes`
+        and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+        to maintain the information.
+
         Examples
         --------
 
@@ -632,6 +648,14 @@ class DGLHeteroGraph(object):
             and ``edata`` of the resulting graph under name ``dgl.NID`` and ``dgl.EID``,
             respectively.
 
+        Notes
+        -----
+
+        This function discards the batch information. Please use
+        :func:`dgl.DGLGraph.set_batch_num_nodes`
+        and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+        to maintain the information.
+
         Examples
         --------
 
@@ -1296,8 +1320,8 @@ class DGLHeteroGraph(object):
 
         Manually set batch information
 
-        >>> g.set_batch_num_nodes(torch.tensor([3, 3])
-        >>> g.set_batch_num_edges(torch.tensor([3, 3])
+        >>> g.set_batch_num_nodes(torch.tensor([3, 3]))
+        >>> g.set_batch_num_edges(torch.tensor([3, 3]))
 
         Unbatch the graph.
 
@@ -1441,17 +1465,17 @@ class DGLHeteroGraph(object):
 
         Manually set batch information
 
-        >>> g.set_batch_num_nodes(torch.tensor([3, 3])
-        >>> g.set_batch_num_edges(torch.tensor([3, 3])
+        >>> g.set_batch_num_nodes(torch.tensor([3, 3]))
+        >>> g.set_batch_num_edges(torch.tensor([3, 3]))
 
         Unbatch the graph.
 
         >>> dgl.unbatch(g)
-            [Graph(num_nodes=3, num_edges=3,
-                  ndata_schemes={}
-                  edata_schemes={}), Graph(num_nodes=3, num_edges=3,
-                  ndata_schemes={}
-                  edata_schemes={})]
+        [Graph(num_nodes=3, num_edges=3,
+              ndata_schemes={}
+              edata_schemes={}), Graph(num_nodes=3, num_edges=3,
+              ndata_schemes={}
+              edata_schemes={})]
 
         Create a heterogeneous graph.
 

python/dgl/subgraph.py

@@ -59,6 +59,14 @@ def node_subgraph(graph, nodes, store_ids=True):
     G : DGLGraph
         The subgraph.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following example uses PyTorch backend.
@@ -187,6 +195,14 @@ def edge_subgraph(graph, edges, preserve_nodes=False, store_ids=True):
     G : DGLGraph
         The subgraph.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following example uses PyTorch backend.
@@ -322,6 +338,14 @@ def in_subgraph(g, nodes):
     DGLGraph
         The subgraph.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following example uses PyTorch backend.
@@ -420,6 +444,14 @@ def out_subgraph(g, nodes):
     DGLGraph
         The subgraph.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following example uses PyTorch backend.
@@ -501,6 +533,14 @@ def node_type_subgraph(graph, ntypes):
     G : DGLGraph
         The subgraph.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following example uses PyTorch backend.
@@ -575,6 +615,14 @@ def edge_type_subgraph(graph, etypes):
     G : DGLGraph
         The subgraph.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following example uses PyTorch backend.

python/dgl/transform.py

@@ -258,6 +258,11 @@ def to_bidirected(g, copy_ndata=False, readonly=None):
     tensors with the input graph. Hence, users should try to avoid in-place operations
     which will be visible to both graphs.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following examples use PyTorch backend.
@@ -354,6 +359,11 @@ def add_reverse_edges(g, readonly=None, copy_ndata=True,
     which will be visible to both graphs. On the contrary, the two graphs do not share
     the same edge feature storage.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     **Homogeneous graphs**
@@ -478,8 +488,11 @@ def line_graph(g, backtracking=True, shared=False):
     * If :attr:`shared` is True, the node features of the resulting graph share the same
       storage with the edge features of the input graph. Hence, users should try to
       avoid in-place operations which will be visible to both graphs.
-
     * The function supports input graph on GPU but copies it to CPU during computation.
+    * This function discards the batch information. Please use
+      :func:`dgl.DGLGraph.set_batch_num_nodes`
+      and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+      to maintain the information.
 
     Examples
     --------
@@ -594,6 +607,11 @@ def khop_graph(g, k, copy_ndata=True):
     tensors with the input graph. Hence, users should try to avoid in-place operations
     which will be visible to both graphs.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
 
@@ -672,6 +690,11 @@ def reverse(g, copy_ndata=True, copy_edata=False, *, share_ndata=None, share_eda
     tensors with the input graph. Hence, users should try to avoid in-place operations
     which will be visible to both graphs.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     **Homogeneous graphs**
@@ -774,6 +797,14 @@ def to_simple_graph(g):
     -------
     DGLGraph
         A simple graph.
+
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
     """
     dgl_warning('dgl.to_simple_graph is renamed to dgl.to_simple in v0.5.')
     return to_simple(g)
@@ -804,6 +835,11 @@ def to_bidirected_stale(g, readonly=True):
     -----
     Please make sure g is a simple graph, otherwise the return value is undefined.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Returns
     -------
     DGLGraph
@@ -896,6 +932,14 @@ def metapath_reachable_graph(g, metapath):
         A homogeneous or unidirectional bipartite graph. It will be on CPU regardless of
         whether the input graph is on CPU or GPU.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     >>> g = dgl.heterograph({
@@ -951,6 +995,10 @@ def add_nodes(g, num, data=None, ntype=None):
       DGL assigns zero features for the newly added nodes.
     * For feature in :attr:`data` but not in :attr:`g`, DGL assigns zero features
       for the existing nodes in the graph.
+    * This function discards the batch information. Please use
+      :func:`dgl.DGLGraph.set_batch_num_nodes`
+      and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+      to maintain the information.
 
     Examples
     --------
@@ -1051,6 +1099,10 @@ def add_edges(g, u, v, data=None, etype=None):
       DGL assigns zero features for the newly added nodes.
     * For feature in :attr:`data` but not in :attr:`g`, DGL assigns zero features
       for the existing nodes in the graph.
+    * This function discards the batch information. Please use
+      :func:`dgl.DGLGraph.set_batch_num_nodes`
+      and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+      to maintain the information.
 
     Examples
     --------
@@ -1148,6 +1200,14 @@ def remove_edges(g, eids, etype=None, store_ids=False):
     DGLGraph
         The graph with edges deleted.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     >>> import dgl
@@ -1214,6 +1274,14 @@ def remove_nodes(g, nids, ntype=None, store_ids=False):
     DGLGraph
         The graph with nodes deleted.
 
+    Notes
+    -----
+
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
 
@@ -1290,6 +1358,10 @@ def add_self_loop(g, etype=None):
       If one wishes to have exactly one self-loop for every node,
       call :func:`remove_self_loop` before invoking :func:`add_self_loop`.
     * Features of the new edges (self-loop edges) will be filled with zeros.
+    * This function discards the batch information. Please use
+      :func:`dgl.DGLGraph.set_batch_num_nodes`
+      and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+      to maintain the information.
 
     Examples
     --------
@@ -1360,6 +1432,11 @@ def remove_self_loop(g, etype=None):
     If a node has multiple self-loops, remove them all. Do nothing for nodes without
     self-loops.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     ---------
 
@@ -1474,6 +1551,11 @@ def compact_graphs(graphs, always_preserve=None, copy_ndata=True, copy_edata=Tru
     tensors with the input graph. Hence, users should try to avoid in-place operations
     which will be visible to both graphs.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     The following code constructs a bipartite graph with 20 users and 10 games, but
@@ -1824,6 +1906,11 @@ def to_simple(g,
     tensors with the input graph. Hence, users should try to avoid in-place operations
     which will be visible to both graphs.
 
+    This function discards the batch information. Please use
+    :func:`dgl.DGLGraph.set_batch_num_nodes`
+    and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
+    to maintain the information.
+
     Examples
     --------
     **Homogeneous Graphs**