message: [tf.distribute] Adding a test for tf.data.Dataset.cache.
diff: @@ -511,6 +511,26 @@ class DistributedIteratorSingleWorkerTest(DistributedIteratorTestBase,
  sess=None,
  split_batch_by=split_batch_by)
  + @combinations.generate(
  + combinations.combine(
  + mode=["eager"],
  + distribution=[
  + strategy_combinations.one_device_strategy,
  + strategy_combinations.mirrored_strategy_with_one_cpu,
  + strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
  + strategy_combinations.tpu_strategy,
  + strategy_combinations.central_storage_strategy_with_two_gpus,
  + ],
  + ))
  + def testCache(self, distribution):
  + dataset = dataset_ops.Dataset.range(10).shuffle(10).cache().batch(1)
  + dist_dataset = distribution.experimental_distribute_dataset(dataset)
  +
  + first_epoch = list(x.numpy() for x in dist_dataset)
  + second_epoch = list(x.numpy() for x in dist_dataset)
  +
  + self.assertEqual(first_epoch, second_epoch)
  +
  class DistributedIteratorTensorTypeTest(DistributedIteratorTestBase,
  parameterized.TestCase):
  @@ -1078,6 +1098,5 @@ class InputTypeSpecTest(test.TestCase, parameterized.TestCase):
  for x in dist_dataset:
  process_inputs(x)
  -
  if __name__ == "__main__":
  test.main()
  