code: def test_set_batch_info(idtype):
      ctx = F.ctx()
  
      g1 = dgl.rand_graph(30, 100).astype(idtype).to(F.ctx())
      g2 = dgl.rand_graph(40, 200).astype(idtype).to(F.ctx())
      bg = dgl.batch([g1, g2])
      batch_num_nodes = F.astype(bg.batch_num_nodes(), idtype)
      batch_num_edges = F.astype(bg.batch_num_edges(), idtype)
  
      # test homogeneous node subgraph
      sg_n = dgl.node_subgraph(
          bg, list(range(10, 20)) + list(range(50, 60))
      )
      induced_nodes = sg_n.ndata["_ID"]
      induced_edges = sg_n.edata["_ID"]
      new_batch_num_nodes = _get_subgraph_batch_info(
          bg.ntypes, [induced_nodes], batch_num_nodes
      )
      new_batch_num_edges = _get_subgraph_batch_info(
          bg.canonical_etypes, [induced_edges], batch_num_edges
      )
      sg_n.set_batch_num_nodes(new_batch_num_nodes)
      sg_n.set_batch_num_edges(new_batch_num_edges)
      subg_n1, subg_n2 = dgl.unbatch(sg_n)
      subg1 = dgl.node_subgraph(g1, list(range(10, 20)))
      subg2 = dgl.node_subgraph(g2, list(range(20, 30)))
      assert subg_n1.num_edges() == subg1.num_edges()
      assert subg_n2.num_edges() == subg2.num_edges()
  
      # test homogeneous edge subgraph
      sg_e = dgl.edge_subgraph(
          bg,
          list(range(40, 70)) + list(range(150, 200)),
          relabel_nodes=False,
      )
      induced_nodes = F.arange(0, bg.num_nodes(), idtype)
      induced_edges = sg_e.edata["_ID"]
      new_batch_num_nodes = _get_subgraph_batch_info(
          bg.ntypes, [induced_nodes], batch_num_nodes
      )
      new_batch_num_edges = _get_subgraph_batch_info(
          bg.canonical_etypes, [induced_edges], batch_num_edges
      )
      sg_e.set_batch_num_nodes(new_batch_num_nodes)
      sg_e.set_batch_num_edges(new_batch_num_edges)
      subg_e1, subg_e2 = dgl.unbatch(sg_e)
      subg1 = dgl.edge_subgraph(
          g1, list(range(40, 70)), relabel_nodes=False
      )
      subg2 = dgl.edge_subgraph(
          g2, list(range(50, 100)), relabel_nodes=False
      )
      assert subg_e1.num_nodes() == subg1.num_nodes()
      assert subg_e2.num_nodes() == subg2.num_nodes()
  
file_path: dgl/tests/python/common/test_batch-graph.py
imports: import unittest
  import backend
  import dgl
  import numpy
  from utils import parametrize_idtype
label_desc: Use image node_id 13 for calling a global function with example usage: test_set_batch_info(idtype) without return types