code: def collect_stats(mod, dataset, chunk_by=-1):
      """Given an annotated graph, create a profile graph to collect profile data from the
      calibration dataset. This pass collects simulated_quantize op input into a tuple.
      Simulated_quantize ops are rewritten to identity mode. The tuple is the output of the profile
      graph.
  
      Parameters
      ----------
      mod: Module
          The simulation graph after annotation.
  
      dataset: Iterable[NDArray]
          The calibration dataset.
  
      chunk_by: optional, int
          The size of chunk to be returned in one iteration. It is meant to be
          used for reducing memory usage. If not specified, return samples for
          all layers in one chunk.
  
      Returns
      -------
      ret: Iterable[list of ndarray]
          List of output data of each layer, chunked by the chunk_by parameter
      """
      logging.info("collecting statistics for calibration...")
      runtime = _get_profile_runtime(mod)
      num_outputs = runtime.get_num_outputs()
      chunk_by = num_outputs if chunk_by == -1 else chunk_by
  
      for i in range(0, num_outputs, chunk_by):
          outputs = [[] for i in range(min(chunk_by, num_outputs - i))]
          for batch in dataset:
              runtime.set_input(**batch)
              runtime.run()
              for j in range(i, min(i + chunk_by, num_outputs)):
                  outputs[j - i].append(runtime.get_output(j).numpy())
          yield [
              np.concatenate(output).reshape(-1) for output in outputs
          ]
  
imports: from __future__ import absolute_import
  import logging
  import multiprocessing
  import numpy
  import tvm
  import tvm.driver
  from tvm.ir import IRModule
  from .None import _quantize
  from .None import quantize
  from ..None import op
  from ..None import expr
  from ..None import analysis
  from ..None import build_module
  from ...contrib import graph_executor
  from .kl_divergence import _find_scale_by_kl
label_desc: Use image node_id 2 for calling a global function with example usage: collect_stats(mod, dataset, chunk_by) without return types