class_docstr: Tests for distributed save/load with the new load algorithm.
class_docstr_tok: ["Tests","for","distributed","save\/load","with","the","new","load","algorithm","."]
code: def test_shuffle_chunks(
      self, num_elements: int, num_repetitions: int
  ):
      cluster = data_service_test_base.TestCluster(
          num_workers=5, snapshot_max_chunk_size_bytes=8
      )
      snapshot_dir = data_service_test_base.TempDir()
      dataset = dataset_ops.Dataset.range(num_elements)
      self.evaluate(
          distributed_save_op.distributed_save(
              dataset,
              snapshot_dir.full_path,
              cluster.dispatcher_address(),
          )
      )
  
      def _interleave_shuffled_chunks(
          datasets: dataset_ops.Dataset,
      ) -> dataset_ops.Dataset:
          # Does not shuffle when the snapshot is being written (first repetition).
          # Otherwise, loading will be blocked to fill the shuffle buffer, not to
          # read the partial chunks.
          datasets = cond.cond(
              math_ops.greater(datasets.cardinality(), 0),
              lambda: datasets.shuffle(
                  buffer_size=datasets.cardinality()
              ),
              lambda: datasets,
          )
          return datasets.interleave(
              lambda x: x,
              cycle_length=multiprocessing.cpu_count(),
              num_parallel_calls=dataset_ops.AUTOTUNE,
          )
  
      dataset = dataset_ops.Dataset.range(num_repetitions).flat_map(
          lambda _: load_op._load_with_retry(
              snapshot_dir.full_path,
              reader_func=_interleave_shuffled_chunks,
          )
      )
  
      output = self.getDatasetOutput(dataset)
      output_per_repetition = [
          output[i : i + num_elements]
          for i in range(0, len(output), num_elements)
      ]
      self.assertLen(output_per_repetition, num_repetitions)
      for i in range(1, num_repetitions):
          self.assertNotEqual(
              output_per_repetition[i], output_per_repetition[i - 1]
          )
  
imports: import multiprocessing
  import os
  import threading
  import time
  from typing import Callable, Optional
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.data.experimental.kernel_tests.service import test_base
  from tensorflow.python.data.experimental.ops import data_service_ops
  from tensorflow.python.data.experimental.ops import distributed_save_op
  from tensorflow.python.data.kernel_tests import checkpoint_test_base
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import load_op
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import errors
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import cond
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 4 for calling the DistributedSaveLoadTest obj's underlying member method code with example usage: obj.test_shuffle_chunks(num_elements, num_repetitions) and returns: datasets