message: [tf.data] Adding a test for chaining optimizations.
diff: @@ -208,6 +208,26 @@ class OptimizeDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  dataset = dataset.with_options(options)
  self.assertDatasetProduces(dataset, expected_output=[[0]])
  + @combinations.generate(test_base.default_test_combinations())
  + def testOptimizationDoubleOptimizeDatasetNested(self):
  + def flat_map_fn(_):
  + dataset = dataset_ops.Dataset.from_tensors(0)
  + dataset = dataset.apply(testing.assert_next(["MapAndBatch"]))
  + dataset = dataset.skip(0)
  + # Should be fused by map and batch fusion
  + dataset = dataset.map(lambda x: x)
  + dataset = dataset.batch(1)
  + return dataset
  +
  + dataset = dataset_ops.Dataset.from_tensors(0)
  + dataset = dataset.flat_map(flat_map_fn)
  + dataset = dataset_ops._OptimizeDataset(dataset, ["map_and_batch_fusion"],
  + [], [])
  + dataset = dataset_ops._OptimizeDataset(dataset, ["noop_elimination"], [],
  + [])
  +
  + self.assertDatasetProduces(dataset, expected_output=[[0]])
  +
  @combinations.generate(
  combinations.times(
  test_base.default_test_combinations(),
  @@ -542,5 +562,6 @@ class OptimizeDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  self.assertEqual(cpu_budget, 1000)
  self.assertEqual(ram_budget, 999999999)
  +
  if __name__ == "__main__":
  test.main()
  