message: When rebatching a dataset, add a prefetch if there was one
  Summary: Ref T23733
  Test Plan: All tests pass
  Reviewers: #tensorflow, davidn, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved
  Maniphest Tasks: T23733
diff: @@ -24,6 +24,7 @@ import weakref
  from tensorflow.python.data.experimental.ops import cardinality
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import distribution_strategy_context
  +from tensorflow.python.distribute import input_lib
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import device as tf_device
  from tensorflow.python.framework import ops
  @@ -98,9 +99,14 @@ def _get_dataset_and_count(x, y, batch_size):
  size = adapter.get_size()
  if adapter.has_partial_batch():
  + _, _, prefetch_buffer = input_lib._get_dataset_attributes(dataset) # pylint: disable=protected-access
  dataset = dataset.unbatch()
  # Remove the partial batch from the dataset.
  dataset = dataset.batch(batch_size, drop_remainder=True)
  +
  + if prefetch_buffer is not None:
  + dataset = dataset.prefetch(prefetch_buffer)
  +
  size -= 1
  return dataset, size
  