code: def unbatch(g, node_split=None, edge_split=None):
      """Revert the batch operation by split the given graph into a list of small ones.
  
      This is the reverse operation of :func:``dgl.batch``. If the ``node_split``
      or the ``edge_split`` is not given, it calls :func:`DGLGraph.batch_num_nodes`
      and :func:`DGLGraph.batch_num_edges` of the input graph to get the information.
  
      If the ``node_split`` or the ``edge_split`` arguments are given,
      it will partition the graph according to the given segments. One must assure
      that the partition is valid -- edges of the i^th graph only connect nodes
      belong to the i^th graph. Otherwise, DGL will throw an error.
  
      The function supports heterograph input, in which case the two split
      section arguments shall be of dictionary type -- similar to the
      :func:`DGLGraph.batch_num_nodes`
      and :func:`DGLGraph.batch_num_edges` attributes of a heterograph.
  
      Parameters
      ----------
      g : DGLGraph
          Input graph to unbatch.
      node_split : Tensor, dict[str, Tensor], optional
          Number of nodes of each result graph.
      edge_split : Tensor, dict[str, Tensor], optional
          Number of edges of each result graph.
  
      Returns
      -------
      list[DGLGraph]
          Unbatched list of graphs.
  
      Examples
      --------
  
      Unbatch a batched graph
  
      >>> 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])))
      >>> # add features
      >>> 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])
      >>> f1, f2 = dgl.unbatch(bg)
      >>> f1
      Graph(num_nodes=4, num_edges=3,
            ndata_schemes={‘x’ : Scheme(shape=(3,), dtype=torch.float32)}
            edata_schemes={‘w’ : Scheme(shape=(2,), dtype=torch.float32)})
      >>> f2
      Graph(num_nodes=3, num_edges=4,
            ndata_schemes={‘x’ : Scheme(shape=(3,), dtype=torch.float32)}
            edata_schemes={‘w’ : Scheme(shape=(2,), dtype=torch.float32)})
  
      With provided split arguments:
  
      >>> g1 = dgl.graph((th.tensor([0, 1, 2]), th.tensor([1, 2, 3])))
      >>> g2 = dgl.graph((th.tensor([0, 0, 0, 1]), th.tensor([0, 1, 2, 0])))
      >>> g3 = dgl.graph((th.tensor([0]), th.tensor([1])))
      >>> bg = dgl.batch([g1, g2, g3])
      >>> bg.batch_num_nodes()
      tensor([4, 3, 2])
      >>> bg.batch_num_edges()
      tensor([3, 4, 1])
      >>> # unbatch but merge g2 and g3
      >>> f1, f2 = dgl.unbatch(bg, th.tensor([4, 5]), th.tensor([3, 5]))
      >>> f1
      Graph(num_nodes=4, num_edges=3,
            ndata_schemes={}
            edata_schemes={})
      >>> f2
      Graph(num_nodes=5, num_edges=5,
            ndata_schemes={}
            edata_schemes={})
  
      Heterograph input
  
      >>> 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])
      >>> f1, f2 = dgl.unbatch(bhg)
      >>> f1
      Graph(num_nodes={'user': 2, 'game': 1},
            num_edges={('user', 'plays', 'game'): 2},
            metagraph=[('drug', 'game')])
      >>> f2
      Graph(num_nodes={'user': 1, 'game': 3},
            num_edges={('user', 'plays', 'game'): 3},
            metagraph=[('drug', 'game')])
  
      See Also
      --------
      batch
      """
      num_split = None
      # Parse node_split
      if node_split is None:
          node_split = {
              ntype: g.batch_num_nodes(ntype) for ntype in g.ntypes
          }
      elif not isinstance(node_split, Mapping):
          if len(g.ntypes) != 1:
              raise DGLError(
                  "Must provide a dictionary for argument node_split when"
                  " there are multiple node types."
              )
          node_split = {g.ntypes[0]: node_split}
      if node_split.keys() != set(g.ntypes):
          raise DGLError("Must specify node_split for each node type.")
      for split in node_split.values():
          if num_split is not None and num_split != len(split):
              raise DGLError(
                  "All node_split and edge_split must specify the same number"
                  " of split sizes."
              )
          num_split = len(split)
  
      # Parse edge_split
      if edge_split is None:
          edge_split = {
              etype: g.batch_num_edges(etype)
              for etype in g.canonical_etypes
          }
      elif not isinstance(edge_split, Mapping):
          if len(g.etypes) != 1:
              raise DGLError(
                  "Must provide a dictionary for argument edge_split when"
                  " there are multiple edge types."
              )
          edge_split = {g.canonical_etypes[0]: edge_split}
      if edge_split.keys() != set(g.canonical_etypes):
          raise DGLError(
              "Must specify edge_split for each canonical edge type."
          )
      for split in edge_split.values():
          if num_split is not None and num_split != len(split):
              raise DGLError(
                  "All edge_split and edge_split must specify the same number"
                  " of split sizes."
              )
          num_split = len(split)
  
      node_split = {
          k: F.asnumpy(split).tolist()
          for k, split in node_split.items()
      }
      edge_split = {
          k: F.asnumpy(split).tolist()
          for k, split in edge_split.items()
      }
  
      # Split edges for each relation
      edge_dict_per = [{} for i in range(num_split)]
      for rel in g.canonical_etypes:
          srctype, etype, dsttype = rel
          srcnid_off = dstnid_off = 0
          u, v = g.edges(order="eid", etype=rel)
          us = F.split(u, edge_split[rel], 0)
          vs = F.split(v, edge_split[rel], 0)
          for i, (subu, subv) in enumerate(zip(us, vs)):
              edge_dict_per[i][rel] = (
                  subu - srcnid_off,
                  subv - dstnid_off,
              )
              srcnid_off += node_split[srctype][i]
              dstnid_off += node_split[dsttype][i]
      num_nodes_dict_per = [
          {k: split[i] for k, split in node_split.items()}
          for i in range(num_split)
      ]
  
      # Create graphs
      gs = [
          convert.heterograph(
              edge_dict, num_nodes_dict, idtype=g.idtype
          )
          for edge_dict, num_nodes_dict in zip(
              edge_dict_per, num_nodes_dict_per
          )
      ]
  
      # Unbatch node features
      for ntype in g.ntypes:
          for key, feat in g.nodes[ntype].data.items():
              subfeats = F.split(feat, node_split[ntype], 0)
              for subg, subf in zip(gs, subfeats):
                  subg.nodes[ntype].data[key] = subf
  
      # Unbatch edge features
      for etype in g.canonical_etypes:
          for key, feat in g.edges[etype].data.items():
              subfeats = F.split(feat, edge_split[etype], 0)
              for subg, subf in zip(gs, subfeats):
                  subg.edges[etype].data[key] = subf
  
      return gs
  
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 3 for calling a global function with example usage: unbatch(g, node_split, edge_split) and returns: gs