code: def _get_subgraph_batch_info(
      keys, induced_indices_arr, batch_num_objs
  ):
      """Internal function to compute batch information for subgraphs.
      Parameters
      ----------
      keys : List[str]
          The node/edge type keys.
      induced_indices_arr : List[Tensor]
          The induced node/edge index tensor for all node/edge types.
      batch_num_objs : Tensor
          Number of nodes/edges for each graph in the original batch.
      Returns
      -------
      Mapping[str, Tensor]
          A dictionary mapping all node/edge type keys to the ``batch_num_objs``
          array of corresponding graph.
      """
      bucket_offset = np.expand_dims(
          np.cumsum(F.asnumpy(batch_num_objs), 0), -1
      )  # (num_bkts, 1)
      ret = {}
      for key, induced_indices in zip(keys, induced_indices_arr):
          # NOTE(Zihao): this implementation is not efficient and we can replace it with
          # binary search in the future.
          induced_indices = np.expand_dims(
              F.asnumpy(induced_indices), 0
          )  # (1, num_nodes)
          new_offset = np.sum(
              (induced_indices < bucket_offset), 1
          )  # (num_bkts,)
          # start_offset = [0] + [new_offset[i-1] for i in range(1, n_bkts)]
          start_offset = np.concatenate(
              [np.zeros((1,)), new_offset[:-1]], 0
          )
          new_batch_num_objs = new_offset - start_offset
          ret[key] = F.tensor(
              new_batch_num_objs, dtype=F.dtype(batch_num_objs)
          )
      return ret
  
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 12 for calling a global function with example usage: _get_subgraph_batch_info(keys, induced_indices_arr, batch_num_objs) and returns: ret