[XLA:GPU][TileAnalysis][NFC] Fix order of function composition in com…

…ments. ∘ :: (b -> c) -> (a -> b) -> (a -> c) FUTURE_COPYBARA_INTEGRATE_REVIEW=#62782 from DanielYang59:warn-shuffle 134fc24 PiperOrigin-RevId: 615527130
tensorflow · Mar 13, 2024 · 1ec7ccf · 1ec7ccf
1 parent a71dbba
commit 1ec7ccf
Show file tree

Hide file tree

Showing 15 changed files with 52 additions and 25 deletions.
diff --git a/tensorflow/core/data/service/snapshot/BUILD b/tensorflow/core/data/service/snapshot/BUILD
@@ -127,6 +127,7 @@ cc_library(
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:path",
+        "@local_tsl//tsl/platform:random",
         "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/profiler/lib:traceme",
     ],

diff --git a/tensorflow/core/data/service/snapshot/parallel_tfrecord_writer.cc b/tensorflow/core/data/service/snapshot/parallel_tfrecord_writer.cc
@@ -35,6 +35,7 @@ limitations under the License.
 #include "tsl/platform/env.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/path.h"
+#include "tsl/platform/random.h"
 #include "tsl/platform/statusor.h"
 #include "tsl/platform/threadpool.h"
 #include "tsl/profiler/lib/traceme.h"
@@ -164,7 +165,7 @@ ParallelTFRecordWriter::GetNextRecord(const std::string& filename)
   std::vector<Tensor> record = std::move(buffer_.front());
   ByteSize estimated_size = EstimatedSize(record);
   LOG_EVERY_N_SEC(INFO, 1) << "Writing TFRecord of " << estimated_size
-                           << " to file " << file_prefix_ << "*.";
+                           << " to file " << filename << "*.";
   ++file_stats_[filename].num_records;
   file_stats_[filename].estimated_size += estimated_size;
   buffer_.pop_front();
@@ -188,7 +189,8 @@ absl::Status ParallelTFRecordWriter::DeleteEmptyFile(
 }
 
 absl::StatusOr<std::string> ParallelTFRecordWriter::GetUniqueFile() const {
-  std::string filename = absl::StrCat(file_prefix_, "__shard__");
+  std::string filename = absl::StrCat(file_prefix_, "__shard__",
+                                      absl::Hex(tsl::random::New64()), "_");
   if (!env_->CreateUniqueFileName(&filename, ".tfrecord")) {
     return absl::InternalError(
         absl::StrCat("Failed to write file ", filename,

diff --git a/tensorflow/core/tfrt/saved_model/saved_model.cc b/tensorflow/core/tfrt/saved_model/saved_model.cc
@@ -111,6 +111,12 @@ auto* lazy_loading_count = monitoring::Counter<3>::New(
     "/tensorflow/tfrt/lazy_loading_count", "The total number of lazy loadings.",
     "model_name", "model_version", "use_graph_executor");
 
+auto* saved_model_graph_executor_mode = monitoring::Counter<3>::New(
+    "/tensorflow/tfrt/saved_model/graph_executor_mode",
+    "Record the total number of imported savedmodel using different graph "
+    "executor modes (BEF vs MLRT interpreter)",
+    "model_name", "model_version", "mode");
+
 auto* saved_model_import_time_seconds =
     tensorflow::monitoring::Gauge<int64_t, 1>::New(
         "/tensorflow/tfrt/saved_model/import_time",
@@ -603,11 +609,25 @@ SavedModelImpl::LoadSavedModel(Options options,
     tensorflow::tf_mlrt::RegisterTfMlrtBatchKernels(*kernel_registry);
 
     if (options.graph_execution_options.enable_mlrt) {
+      saved_model_graph_executor_mode
+          ->GetCell(
+              options.graph_execution_options.model_metadata.name(),
+              absl::StrCat(
+                  options.graph_execution_options.model_metadata.version()),
+              "mlrt")
+          ->IncrementBy(1);
       ASSIGN_OR_RETURN_IN_COMPILE(
           bytecode, tensorflow::mlrt_compiler::ConvertTfMlirToBytecode(
                         options.graph_execution_options.compile_options,
                         *fallback_state, mlir_module.get(), model_context));
     } else {
+      saved_model_graph_executor_mode
+          ->GetCell(
+              options.graph_execution_options.model_metadata.name(),
+              absl::StrCat(
+                  options.graph_execution_options.model_metadata.version()),
+              "bef")
+          ->IncrementBy(1);
       RETURN_IF_ERROR_IN_COMPILE(tensorflow::ConvertTfMlirToBef(
           options.graph_execution_options.compile_options, mlir_module.get(),
           &bef, model_context, fallback_state.get()));

diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py
@@ -1408,7 +1408,7 @@ def enumerate(self, start=0, name=None) -> "DatasetV2":
     return Dataset.zip((range_dataset, self), name=name)
 
   def shuffle(
-      self, buffer_size, seed=None, reshuffle_each_iteration=None, name=None
+      self, buffer_size, seed=None, reshuffle_each_iteration=True, name=None
   ) -> "DatasetV2":
     """Randomly shuffles the elements of this dataset.
 
@@ -1424,8 +1424,12 @@ def shuffle(
     maintaining the 1,000 element buffer.
 
     `reshuffle_each_iteration` controls whether the shuffle order should be
-    different for each epoch. In TF 1.X, the idiomatic way to create epochs
-    was through the `repeat` transformation:
+    different for each epoch. However you should avoid using
+    `shuffle(reshuffle_each_iteration=True)`, then `take` and `skip` to split
+    a dataset into training and test sets, which would lead to data leakage (as
+    the entire dataset would be re-shuffled then re-split after each epoch).
+    Please use the `tf.keras.utils.split_dataset` method instead. In TF 1.X,
+    the idiomatic way to create epochs was through the `repeat` transformation:
 
     ```python
     dataset = tf.data.Dataset.range(3)

diff --git a/tensorflow/python/data/ops/shuffle_op.py b/tensorflow/python/data/ops/shuffle_op.py
@@ -26,28 +26,29 @@ def _shuffle(  # pylint: disable=unused-private-name
     input_dataset,
     buffer_size,
     seed=None,
-    reshuffle_each_iteration=None,
-    name=None):
+    reshuffle_each_iteration=True,
+    name=None,
+):
   return _ShuffleDataset(
       input_dataset, buffer_size, seed, reshuffle_each_iteration, name=name)
 
 
 class _ShuffleDataset(dataset_ops.UnaryUnchangedStructureDataset):
   """A `Dataset` that randomly shuffles the elements of its input."""
 
-  def __init__(self,
-               input_dataset,
-               buffer_size,
-               seed=None,
-               reshuffle_each_iteration=None,
-               name=None):
+  def __init__(
+      self,
+      input_dataset,
+      buffer_size,
+      seed=None,
+      reshuffle_each_iteration=True,
+      name=None,
+  ):
     """See `Dataset.shuffle()` for details."""
     self._input_dataset = input_dataset
     self._buffer_size = ops.convert_to_tensor(
         buffer_size, dtype=dtypes.int64, name="buffer_size")
     self._seed, self._seed2 = random_seed.get_seed(seed)
-    if reshuffle_each_iteration is None:
-      reshuffle_each_iteration = True
     self._reshuffle_each_iteration = reshuffle_each_iteration
     self._name = name
 

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt
@@ -160,7 +160,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-fixed-length-record-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-fixed-length-record-dataset.pbtxt
@@ -162,7 +162,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-t-f-record-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-t-f-record-dataset.pbtxt
@@ -161,7 +161,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-text-line-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-text-line-dataset.pbtxt
@@ -162,7 +162,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt
@@ -162,7 +162,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.pbtxt
@@ -163,7 +163,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.pbtxt
@@ -162,7 +162,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/tensorflow/tools/api/golden/v2/tensorflow.experimental.dtensor.-d-tensor-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.experimental.dtensor.-d-tensor-dataset.pbtxt
@@ -163,7 +163,7 @@ tf_class {
   }
   member_method {
     name: "shuffle"
-    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "skip"

diff --git a/third_party/xla/xla/service/gpu/model/indexing_map.cc b/third_party/xla/xla/service/gpu/model/indexing_map.cc
@@ -999,7 +999,6 @@ IndexingMap ComposeIndexingMaps(const IndexingMap& first,
     return IndexingMap::GetUndefined();
   }
   AffineMap producer_affine_map = second.GetAffineMap();
-  // map1.compose(map2) computes map2 ∘ map1 for some reason.
   AffineMap composed_map = producer_affine_map.compose(first.GetAffineMap());
 
   // The symbols in the composed map, i.e. combined

diff --git a/third_party/xla/xla/service/gpu/model/indexing_map.h b/third_party/xla/xla/service/gpu/model/indexing_map.h
@@ -272,7 +272,7 @@ std::ostream& operator<<(std::ostream& out, const IndexingMap& indexing_map);
 bool operator==(const IndexingMap& lhs, const IndexingMap& rhs);
 IndexingMap operator*(const IndexingMap& lhs, const IndexingMap& rhs);
 
-// Composes affine maps, i.e. first ∘ second.
+// Composes affine maps, i.e. second ∘ first.
 IndexingMap ComposeIndexingMaps(const IndexingMap& first,
                                 const IndexingMap& second);