code: def batch(graphs, ndata=ALL, edata=ALL):
      r"""Batch a collection of :class:`DGLGraph` s into one graph for more efficient
      graph computation.
  
      Each input graph becomes one disjoint component of the batched graph. The nodes
      and edges are relabeled to be disjoint segments:
  
      =================  =========  =================  ===  =========
                         graphs[0]  graphs[1]          ...  graphs[k]
      =================  =========  =================  ===  =========
      Original node ID   0 ~ N_0    0 ~ N_1            ...  0 ~ N_k
      New node ID        0 ~ N_0    N_0 ~ N_0+N_1      ...  \sum_{i=0}^{k-1} N_i ~
                                                            \sum_{i=0}^k N_i
      =================  =========  =================  ===  =========
  
      Because of this, many of the computations on a batched graph are the same as if
      performed on each graph individually, but become much more efficient
      since they can be parallelized easily. This makes ``dgl.batch`` very useful
      for tasks dealing with many graph samples such as graph classification tasks.
  
      For heterograph inputs, they must share the same set of relations (i.e., node types
      and edge types) and the function will perform batching on each relation one by one.
      Thus, the result is also a heterograph and has the same set of relations as the inputs.
  
      The numbers of nodes and edges of the input graphs are accessible via the
      :func:`DGLGraph.batch_num_nodes` and :func:`DGLGraph.batch_num_edges` attributes
      of the resulting graph. For homogeneous graphs, they are 1D integer tensors,
      with each element being the number of nodes/edges of the corresponding input graph. For
      heterographs, they are dictionaries of 1D integer tensors, with node
      type or edge type as the keys.
  
      The function supports batching batched graphs. The batch size of the result
      graph is the sum of the batch sizes of all the input graphs.
  
      By default, node/edge features are batched by concatenating the feature tensors
      of all input graphs. This thus requires features of the same name to have
      the same data type and feature size. One can pass ``None`` to the ``ndata``
      or ``edata`` argument to prevent feature batching, or pass a list of strings
      to specify which features to batch.
  
      To unbatch the graph back to a list, use the :func:`dgl.unbatch` function.
  
      Parameters
      ----------
      graphs : list[DGLGraph]
          Input graphs.
      ndata : list[str], None, optional
          Node features to batch.
      edata : list[str], None, optional
          Edge features to batch.
  
      Returns
      -------
      DGLGraph
          Batched graph.
  
      Examples
      --------
  
      Batch homogeneous graphs
  
      >>> import dgl
      >>> import torch as th
      >>> # 4 nodes, 3 edges
      >>> g1 = dgl.graph((th.tensor([0, 1, 2]), th.tensor([1, 2, 3])))
      >>> # 3 nodes, 4 edges
      >>> g2 = dgl.graph((th.tensor([0, 0, 0, 1]), th.tensor([0, 1, 2, 0])))
      >>> bg = dgl.batch([g1, g2])
      >>> bg
      Graph(num_nodes=7, num_edges=7,
            ndata_schemes={}
            edata_schemes={})
      >>> bg.batch_size
      2
      >>> bg.batch_num_nodes()
      tensor([4, 3])
      >>> bg.batch_num_edges()
      tensor([3, 4])
      >>> bg.edges()
      (tensor([0, 1, 2, 4, 4, 4, 5], tensor([1, 2, 3, 4, 5, 6, 4]))
  
      Batch batched graphs
  
      >>> bbg = dgl.batch([bg, bg])
      >>> bbg.batch_size
      4
      >>> bbg.batch_num_nodes()
      tensor([4, 3, 4, 3])
      >>> bbg.batch_num_edges()
      tensor([3, 4, 3, 4])
  
      Batch graphs with feature data
  
      >>> g1.ndata['x'] = th.zeros(g1.num_nodes(), 3)
      >>> g1.edata['w'] = th.ones(g1.num_edges(), 2)
      >>> g2.ndata['x'] = th.ones(g2.num_nodes(), 3)
      >>> g2.edata['w'] = th.zeros(g2.num_edges(), 2)
      >>> bg = dgl.batch([g1, g2])
      >>> bg.ndata['x']
      tensor([[0, 0, 0],
              [0, 0, 0],
              [0, 0, 0],
              [0, 0, 0],
              [1, 1, 1],
              [1, 1, 1],
              [1, 1, 1]])
      >>> bg.edata['w']
      tensor([[1, 1],
              [1, 1],
              [1, 1],
              [0, 0],
              [0, 0],
              [0, 0],
              [0, 0]])
  
      Batch heterographs
  
      >>> hg1 = dgl.heterograph({
      ...     ('user', 'plays', 'game') : (th.tensor([0, 1]), th.tensor([0, 0]))})
      >>> hg2 = dgl.heterograph({
      ...     ('user', 'plays', 'game') : (th.tensor([0, 0, 0]), th.tensor([1, 0, 2]))})
      >>> bhg = dgl.batch([hg1, hg2])
      >>> bhg
      Graph(num_nodes={'user': 3, 'game': 4},
            num_edges={('user', 'plays', 'game'): 5},
            metagraph=[('drug', 'game')])
      >>> bhg.batch_size
      2
      >>> bhg.batch_num_nodes()
      {'user' : tensor([2, 1]), 'game' : tensor([1, 3])}
      >>> bhg.batch_num_edges()
      {('user', 'plays', 'game') : tensor([2, 3])}
  
      See Also
      --------
      unbatch
      """
      if len(graphs) == 0:
          raise DGLError("The input list of graphs cannot be empty.")
      if not (
          is_all(ndata) or isinstance(ndata, list) or ndata is None
      ):
          raise DGLError(
              "Invalid argument ndata: must be a string list but got {}.".format(
                  type(ndata)
              )
          )
      if not (
          is_all(edata) or isinstance(edata, list) or edata is None
      ):
          raise DGLError(
              "Invalid argument edata: must be a string list but got {}.".format(
                  type(edata)
              )
          )
      if any(g.is_block for g in graphs):
          raise DGLError("Batching a MFG is not supported.")
  
      relations = list(graphs[0].canonical_etypes)
      relation_ids = [graphs[0].get_etype_id(r) for r in relations]
      ntypes = list(graphs[0].ntypes)
      ntype_ids = [graphs[0].get_ntype_id(n) for n in ntypes]
      etypes = [etype for _, etype, _ in relations]
  
      gidx = disjoint_union(
          graphs[0]._graph.metagraph, [g._graph for g in graphs]
      )
      retg = DGLGraph(gidx, ntypes, etypes)
  
      # Compute batch num nodes
      bnn = {}
      for ntype in ntypes:
          bnn[ntype] = F.cat(
              [g.batch_num_nodes(ntype) for g in graphs], 0
          )
      retg.set_batch_num_nodes(bnn)
  
      # Compute batch num edges
      bne = {}
      for etype in relations:
          bne[etype] = F.cat(
              [g.batch_num_edges(etype) for g in graphs], 0
          )
      retg.set_batch_num_edges(bne)
  
      # Batch node feature
      if ndata is not None:
          for ntype_id, ntype in zip(ntype_ids, ntypes):
              all_empty = all(
                  g._graph.num_nodes(ntype_id) == 0 for g in graphs
              )
              frames = [
                  g._node_frames[ntype_id]
                  for g in graphs
                  if g._graph.num_nodes(ntype_id) > 0 or all_empty
              ]
              # TODO: do we require graphs with no nodes/edges to have the same schema?  Currently
              # we allow empty graphs to have no features during batching.
              ret_feat = _batch_feat_dicts(
                  frames, ndata, 'nodes["{}"].data'.format(ntype)
              )
              retg.nodes[ntype].data.update(ret_feat)
  
      # Batch edge feature
      if edata is not None:
          for etype_id, etype in zip(relation_ids, relations):
              all_empty = all(
                  g._graph.num_edges(etype_id) == 0 for g in graphs
              )
              frames = [
                  g._edge_frames[etype_id]
                  for g in graphs
                  if g._graph.num_edges(etype_id) > 0 or all_empty
              ]
              # TODO: do we require graphs with no nodes/edges to have the same schema?  Currently
              # we allow empty graphs to have no features during batching.
              ret_feat = _batch_feat_dicts(
                  frames, edata, "edges[{}].data".format(etype)
              )
              retg.edges[etype].data.update(ret_feat)
  
      return retg
  
file_path: dgl/python/dgl/batch.py
imports: from collections.abc import Mapping
  from .None import backend, convert, utils
  from .base import ALL, DGLError, EID, is_all, NID
  from .heterograph import DGLGraph
  from .heterograph_index import disjoint_union, slice_gidx
label_desc: Use image node_id 1 for calling a global function with example usage: batch(graphs, ndata, edata) and returns: retg