message: Optimize the dataset for dataset benchmark by default
  Summary: Fixes T18116
  Test Plan: all tests pass
  Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, davidn
  Maniphest Tasks: T18116
diff: @@ -25,7 +25,8 @@ from tensorflow.python.data.ops import dataset_ops
  def dataset_benchmark(dataset,
  number_of_epochs,
  elements_per_epochs,
  - print_stats=True):
  + print_stats=True,
  + apply_options=True):
  """Allows the user to benchmark performance of a `tf.data.Dataset`.
  Args:
  @@ -34,6 +35,8 @@ def dataset_benchmark(dataset,
  elements_per_epochs: The number of elements there are in each epoch.
  print_stats: Whether to print statistics about the performance to the
  console.
  + apply_options: Whether to apply optimization options which can improve the
  + dataset performance.
  Returns:
  A JSON string with performance statistics, which records the following
  @@ -62,6 +65,9 @@ def dataset_benchmark(dataset,
  "`tf.data.Dataset`, but got %s "
  "instead." % (str(dataset)))
  + if apply_options:
  + dataset = dataset._apply_options() # pylint: disable=protected-access
  +
  try:
  dataset_variant = dataset._variant_tensor # pylint: disable=protected-access
  except TypeError:
  @@ -110,8 +116,11 @@ def infeed_benchmark(infeed_queue,
  return TypeError("Expected `infeed_queue` argument to be of type "
  "`ipu_infeed_queue.IPUInfeedQueue`, but got %s "
  "instead." % (str(infeed_queue)))
  + # Don't need to apply options because the infeed queue already applies them.
  + apply_options = False
  return dataset_benchmark(
  infeed_queue._dataset, # pylint: disable=protected-access
  number_of_epochs,
  elements_per_epochs,
  - print_stats=print_stats)
  + print_stats=print_stats,
  + apply_options=apply_options)
  