code: def _initialize_num_chunks(g, num_chunks, kwargs=None):
      """Initialize num_chunks for each node/edge.
  
      Parameters
      ----------
      g: DGLGraph
          Graph to be chunked.
      num_chunks: int
          Default number of chunks to be applied onto node/edge data.
      kwargs: dict
          Key word arguments to specify details for each node/edge data.
  
      Returns
      -------
      num_chunks_data: dict
          Detailed number of chunks for each node/edge.
      """
  
      def _init(g, num_chunks, key, kwargs=None):
          chunks_data = kwargs.get(key, None)
          is_node = "_node" in key
          data_types = g.ntypes if is_node else g.canonical_etypes
          if isinstance(chunks_data, int):
              chunks_data = {
                  data_type: chunks_data for data_type in data_types
              }
          elif isinstance(chunks_data, dict):
              for data_type in data_types:
                  if data_type not in chunks_data:
                      chunks_data[data_type] = num_chunks
          else:
              chunks_data = {
                  data_type: num_chunks for data_type in data_types
              }
          for _, data in chunks_data.items():
              if isinstance(data, dict):
                  n_chunks = list(data.values())
              else:
                  n_chunks = [data]
              assert all(
                  isinstance(v, int) for v in n_chunks
              ), "num_chunks for each data type should be int."
          return chunks_data
  
      num_chunks_data = {}
      for key in [
          "num_chunks_nodes",
          "num_chunks_edges",
          "num_chunks_node_data",
          "num_chunks_edge_data",
      ]:
          num_chunks_data[key] = _init(
              g, num_chunks, key, kwargs=kwargs
          )
      return num_chunks_data
  
file_path: dgl/tests/tools/pytest_utils.py
imports: import json
  import logging
  import os
  import dgl
  import numpy
  import torch
  from distpartitioning import array_readwriter
  from distpartitioning.array_readwriter.parquet import ParquetArrayParser
  from files import setdir
label_desc: Use image node_id 2 for calling a global function with example usage: _initialize_num_chunks(g, num_chunks, kwargs) and returns: num_chunks_data, chunks_data