code: def rebatch_dataset(dataset, num_workers, worker_index):
      from tensorflow.python.data.experimental.ops import distribute
      from tensorflow.python.framework import ops
      from tensorflow.python.ops import control_flow_ops
      from tensorflow.python.ops import math_ops
  
      batch_size = distribute.compute_batch_size(dataset)
  
      def apply_rebatch():
          batch_sizes = distribute.batch_sizes_for_worker(
              batch_size, num_workers, 1, worker_index
          )
          return distribute._RebatchDataset(
              dataset, batch_sizes
          ).prefetch(1)
  
      def apply_legacy_rebatch():
          return distribute._LegacyRebatchDataset(
              dataset, num_workers
          ).prefetch(1)
  
      with ops.colocate_with(dataset._variant_tensor):
          return control_flow_ops.cond(
              math_ops.not_equal(batch_size, -1),
              true_fn=apply_rebatch,
              false_fn=apply_legacy_rebatch,
          )
  
file_path: BigDL/python/nano/src/bigdl/nano/deps/horovod/distributed_utils_horovod.py
imports: import os
  import json
  import shutil
  from tempfile import TemporaryDirectory
  from contextlib import closing
  import socket
  import tensorflow
  from bigdl.nano.utils.common import schedule_processors
  from bigdl.nano.utils.common import invalidInputError
label_desc: Use image node_id 2 for calling a global function with example usage: rebatch_dataset(dataset, num_workers, worker_index) and returns: distribute, distribute, control_flow_ops