[tf.data] Update map parallelization optimization to always default to preserving the order of elements.

PiperOrigin-RevId: 357026486
Change-Id: I86435cb1e5837e3ca27c77ab69968a0a0ce8ad54

Author: jsimsa

tensorflow/core/grappler/optimizers/data/map_parallelization.cc

@@ -32,7 +32,7 @@ namespace grappler {
 namespace {
 
 constexpr char kMapDataset[] = "MapDataset";
-constexpr char kParallelMapDataset[] = "ParallelMapDataset";
+constexpr char kParallelMapDataset[] = "ParallelMapDatasetV2";
 
 NodeDef MakeParallelMap(const string& name, MutableGraphView* graph) {
   // The inputs of the node to be parallelized could be changed by the
@@ -45,8 +45,9 @@ NodeDef MakeParallelMap(const string& name, MutableGraphView* graph) {
                                       &parallel_map);
   parallel_map.set_op(kParallelMapDataset);
   auto* num_parallel_calls = graph_utils::AddScalarConstNode(
-      static_cast<int32>(data::model::kAutotune), graph);
+      static_cast<int64>(data::model::kAutotune), graph);
   parallel_map.add_input(num_parallel_calls->name());
+  AddNodeAttr("deterministic", "true", &parallel_map);
 
   return parallel_map;
 }

tensorflow/core/grappler/optimizers/data/map_parallelization_test.cc

@@ -68,7 +68,7 @@ TEST_P(AutotuneSetting, MapParallelizationTest) {
 
   GraphDef output;
   TF_ASSERT_OK(OptimizeWithMapParallelization(item, &output, autotune));
-  EXPECT_EQ(graph_utils::ContainsNodeWithOp("ParallelMapDataset", output),
+  EXPECT_EQ(graph_utils::ContainsNodeWithOp("ParallelMapDatasetV2", output),
             autotune);
   EXPECT_EQ(graph_utils::ContainsGraphNodeWithName("map", output), !autotune);
 }
@@ -99,7 +99,7 @@ TEST_P(FromFunctionDef, MapParallelizationTest) {
 
   GraphDef output;
   TF_ASSERT_OK(OptimizeWithMapParallelization(item, &output, true));
-  EXPECT_EQ(graph_utils::ContainsNodeWithOp("ParallelMapDataset", output),
+  EXPECT_EQ(graph_utils::ContainsNodeWithOp("ParallelMapDatasetV2", output),
             !from_function_def);
   EXPECT_EQ(graph_utils::ContainsGraphNodeWithName("map", output),
             from_function_def);
@@ -131,7 +131,7 @@ TEST(ParallelizeAssert, MapParallelizationTest) {
 
   GraphDef output;
   TF_ASSERT_OK(OptimizeWithMapParallelization(item, &output, true));
-  EXPECT_TRUE(graph_utils::ContainsNodeWithOp("ParallelMapDataset", output));
+  EXPECT_TRUE(graph_utils::ContainsNodeWithOp("ParallelMapDatasetV2", output));
   EXPECT_TRUE(graph_utils::ContainsGraphNodeWithName("map1", output));
   EXPECT_FALSE(graph_utils::ContainsGraphNodeWithName("map2", output));
 }

tensorflow/python/data/ops/dataset_ops.py

@@ -1808,11 +1808,12 @@ def map(self, map_func, num_parallel_calls=None, deterministic=None):
         If not specified, elements will be processed sequentially. If the value
         `tf.data.AUTOTUNE` is used, then the number of parallel
         calls is set dynamically based on available CPU.
-      deterministic: (Optional.) A boolean controlling whether determinism
-        should be traded for performance by allowing elements to be yielded out
-        of order.  If `deterministic` is `None`, the
-        `tf.data.Options.experimental_deterministic` dataset option (`True` by
-        default) is used to decide whether to run deterministically.
+      deterministic: (Optional.) When `num_parallel_calls` is specified, this
+        boolean controls the order in which the transformation produces
+        elements. If set to `False`, the transformation is allowed to yield
+        elements out of order to trade determinism for performance. If not
+        specified, the `tf.data.Options.experimental_deterministic` option
+        (`True` by default) controls the behavior.
 
     Returns:
       Dataset: A `Dataset`.
@@ -1937,11 +1938,12 @@ def interleave(self,
         from cycle elements synchronously with no parallelism. If the value
         `tf.data.AUTOTUNE` is used, then the number of parallel
         calls is set dynamically based on available CPU.
-      deterministic: (Optional.) A boolean controlling whether determinism
-        should be traded for performance by allowing elements to be produced out
-        of order.  If `deterministic` is `None`, the
-        `tf.data.Options.experimental_deterministic` dataset option (`True` by
-        default) is used to decide whether to run deterministically.
+      deterministic: (Optional.) When `num_parallel_calls` is specified, this
+        boolean controls the order in which the transformation produces
+        elements. If set to `False`, the transformation is allowed to yield
+        elements out of order to trade determinism for performance. If not
+        specified, the `tf.data.Options.experimental_deterministic` option
+        (`True` by default) controls the behavior.
 
     Returns:
       Dataset: A `Dataset`.
@@ -2673,12 +2675,12 @@ def map_with_legacy_function(self,
         If not specified, elements will be processed sequentially. If the value
         `tf.data.AUTOTUNE` is used, then the number of parallel
         calls is set dynamically based on available CPU.
-      deterministic: (Optional.) A boolean controlling whether determinism
-        should be traded for performance by allowing elements to be produced out
-        of order.  If `deterministic` is `None`, the
-        `tf.data.Options.experimental_deterministic` dataset option (`True` by
-        default) is used to decide whether to produce elements
-        deterministically.
+      deterministic: (Optional.) When `num_parallel_calls` is specified, this
+        boolean controls the order in which the transformation produces
+        elements. If set to `False`, the transformation is allowed to yield
+        elements out of order to trade determinism for performance. If not
+        specified, the `tf.data.Options.experimental_deterministic` option
+        (`True` by default) controls the behavior.
 
     Returns:
       Dataset: A `Dataset`.