Merged commit includes the following changes:

467408627 by A. Unique TensorFlower<gardener@tensorflow.org>: Update sqlite version in TF -- 467380418 by A. Unique TensorFlower<gardener@tensorflow.org>: compat: Update forward compatibility horizon to 2022-08-13 -- 467378663 by A. Unique TensorFlower<gardener@tensorflow.org>: Update GraphDef version to 1222. -- 467363891 by A. Unique TensorFlower<gardener@tensorflow.org>: Update TFRT dependency to use revision http://github.com/tensorflow/runtime/commit/b750bc2999cf02abca6ad9eeff0a04ec7bf3b683. -- 467363622 by A. Unique TensorFlower<gardener@tensorflow.org>: [xla:runtime] NFC: Move constraints documentation from jitrt to xla/runtime/constraints -- 467362586 by A. Unique TensorFlower<gardener@tensorflow.org>: [xla:runtime] NFC: Extract JitCompilationContext library from jitrt and move it to xla/runtime -- 467361314 by A. Unique TensorFlower<gardener@tensorflow.org>: Update TFRT dependency to use revision http://github.com/tensorflow/runtime/commit/0a042cbb5275e6ff9a3a7c2748c74df6dcede09e. -- 467360160 by A. Unique TensorFlower<gardener@tensorflow.org>: [xla:runtime] NFC: Extract calling_convention library from jitrt and move it to xla/runtime -- 467341954 by A. Unique TensorFlower<gardener@tensorflow.org>: Op documentation update. update of g3doc/_includes/tf_passes.md -- 467341426 by A. Unique TensorFlower<gardener@tensorflow.org>: Refactor SELECT_V2 in preparation for porting to TFLM. -- 467340678 by A. Unique TensorFlower<gardener@tensorflow.org>: Create some global stat tracking for CompilationEnvironments. This tracking can be used to help debug cases in which multiple CompilationEnvironments are used to compile a single HloModule (which should not happen). -- 467339870 by A. Unique TensorFlower<gardener@tensorflow.org>: Automated rollback of changelist 467224197. 467339756 by A. Unique TensorFlower<gardener@tensorflow.org>: Update TFRT dependency to use revision http://github.com/tensorflow/runtime/commit/b20ec05d272477fa6223213687bb22145df92674. -- 467339529 by A. Unique TensorFlower<gardener@tensorflow.org>: [XLA] Bugfix for gather index parallel partitioning where the sharded non-parallel dims in indices are not handled. -- 467337900 by A. Unique TensorFlower<gardener@tensorflow.org>: [XLA] Minor renamings, refactorings, checks. -- 467337622 by A. Unique TensorFlower<gardener@tensorflow.org>: Remove unneeded dependency. -- 467337170 by A. Unique TensorFlower<gardener@tensorflow.org>: Integrate LLVM at llvm/llvm-project@2c3ca3b684bb Updates LLVM usage to match [2c3ca3b684bb](llvm/llvm-project@2c3ca3b684bb) -- 467335264 by A. Unique TensorFlower<gardener@tensorflow.org>: [SavedModel Fingerprinting] Add hash #5, which represents the checkpoint. The `checkpoint_hash` is a hash of the serialized .index file, which is the metadata file of the TensorBundle containing a string-string table of the name of a tensor to its serialized BundleEntryProto. The BundleEntryProto contains a crc32 hash of the tensor contents, but not the contents of the tensor itself. RFC: tensorflow/community#415 -- 467334010 by A. Unique TensorFlower<gardener@tensorflow.org>: Update TFRT dependency to use revision http://github.com/tensorflow/runtime/commit/76b3fea4cc9d5e7cb8a85798e41a61a55c301578. -- 467332094 by A. Unique TensorFlower<gardener@tensorflow.org>: [xla:runtime] NFC: Extract executable library from jitrt and move it to xla/runtime -- 467324078 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <vector> for 'std::vector' -- 467322782 by A. Unique TensorFlower<gardener@tensorflow.org>: PR #57137: [oneDNN] Skip appending kernel registration to log message for MKL ops Imported from GitHub PR #57137 This PR skips printing kernel registrations for MKL ops since it leads to performance drop for some eager models caused by this commit c04f65d This is a temporary fix and the condition will be removed when support for block format is removed as a more permanent fix. Copybara import of the project: -- 89c4c20 by Kanvi Khanna <kanvi.khanna@intel.com>: [oneDNN] Skip appending kernel registration to log message for MKL ops Merging this change closes #57137 -- 467322425 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <memory> for 'std::unique_ptr' -- 467321561 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. 'int64' is deprecated: Use int64_t instead. -- 467321058 by A. Unique TensorFlower<gardener@tensorflow.org>: PR #57089: [TF-TRT] Adjusting Conv2D Test Tolerance Imported from GitHub PR #57089 This PR adjusts & fixes the unittest tolerance for the test `Conv2DStridedNCHWTest` in INT8 mode. Copybara import of the project: -- 13e4bff by DEKHTIARJonathan <contact@jonathandekhtiar.eu>: [TF-TRT] Adjusting Conv2D Test Tolerance Merging this change closes #57089 -- 467320826 by A. Unique TensorFlower<gardener@tensorflow.org>: PR #55804: [TF-TRT] Various Cleanups & Python Debugging Assertion Improvements Imported from GitHub PR #55804 This PR cleans a few spots in the code base, improves the debuggability of assertion messages in unittests. And replace `distutils.version.LooseVersion` (deprecated) with `packaging.version.Version` (new recommended API). Copybara import of the project: -- a4d15ef by DEKHTIARJonathan <contact@jonathandekhtiar.eu>: [TF-TRT] Various Cleanups & Python Debugging Assertion Improvements Merging this change closes #55804 -- 467320083 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <ostream> for 'std::ostream' -- 467319094 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. 'int64' is deprecated: Use int64_t instead. -- 467318151 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <iterator> for 'std::back_inserter' -- 467316931 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. using decl 'IsSubsetOf' is unused -- 467316097 by A. Unique TensorFlower<gardener@tensorflow.org>: Move passes under tensorflow/compiler/mlir/tensorflow/. -- 467315812 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <memory> for 'std::unique_ptr' -- 467314236 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <memory> for 'std::unique_ptr' -- 467313254 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <vector> for 'std::vector' missing #include <memory> for 'std::make_unique' -- 467312293 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. using decl 'RangeSquareDataset' is unused -- 467311309 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <vector> for 'std::vector' -- 467310637 by A. Unique TensorFlower<gardener@tensorflow.org>: PR #57013: [TF-TRT] Add LogSoftmax Support for TF-TRT Imported from GitHub PR #57013 This PR adds TF-TRT support to `tf.nn.log_softmax` operation. This is performed using the formula `logsoftmax = logits - log(reduce_sum(exp(logits), axis=-1))` . The implemented TRT layers are fused into a single op. @DEKHTIARJonathan @tfeher : Please review the changes. Copybara import of the project: -- 1a8eb9a by Pavani Majety <pmajety@nvidia.com>: Add LogSoftmax conversion Fix Softmax comments [TF-TRT] Move LogSoftmax to use OpConverterBase Fix compiler errors clang-format Undo changes to convert_nodes.cc Fix comments Merging this change closes #57013 -- 467310335 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <array> for 'std::array' -- 467310313 by A. Unique TensorFlower<gardener@tensorflow.org>: Update test config in cross device ops -- 467309032 by A. Unique TensorFlower<gardener@tensorflow.org>: Update TFRT dependency to use revision http://github.com/tensorflow/runtime/commit/eba528ef667653c3554984e5c05573b152c9893b. -- 467308765 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <vector> for 'std::vector' -- 467307702 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <vector> for 'std::vector' -- 467306473 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. 'int64' is deprecated: Use int64_t instead. -- 467306092 by A. Unique TensorFlower<gardener@tensorflow.org>: PR #56771: Add return_index_map argument in ssim() Imported from GitHub PR #56771 Closes #53115 Copybara import of the project: -- 8f5a1b1 by CohenAriel <ariel17112005@gmail.com>: Add return_index_map argument in ssim() Merging this change closes #56771 -- 467305190 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. missing #include <unordered_map> for 'std::unordered_map' missing #include <vector> for 'std::vector' missing #include <memory> for 'std::shared_ptr' -- 467304747 by A. Unique TensorFlower<gardener@tensorflow.org>: [tfrt:jitrt] NFC: Remove Executable::KernelContext It was added before runtime::KernelContext and is not used anywhere. Remove it to avoid confusion. In the future we should reuse runtime::KernelContext as an extension point for user-defined memory allocation etc. -- 467303335 by A. Unique TensorFlower<gardener@tensorflow.org>: #tf-data-service #codehealth Clean up clang-tidy report. 'int64' is deprecated: Use int64_t instead. -- 467301808 by A. Unique TensorFlower<gardener@tensorflow.org>: Changes all local `State` or `TaskState` enum in coordination service into `CoordinatedTaskState` enum in proto. -- 467300580 by A. Unique TensorFlower<gardener@tensorflow.org>: lite: enable variable freezing in tf_tfl_translate tester -- 467298890 by A. Unique TensorFlower<gardener@tensorflow.org>: Update TFRT dependency to use revision http://github.com/tensorflow/runtime/commit/9bb23f7d1ee0e9a55d26c7168790667e5266a74c. -- 467292686 by A. Unique TensorFlower<gardener@tensorflow.org>: [xla:runtime] NFC: Move execution_engine library from jitrt to xla/runtime -- 467280901 by A. Unique TensorFlower<gardener@tensorflow.org>: [GML] Add tests for concat in the GML tiling and fusion pipeline -- 467276349 by A. Unique TensorFlower<gardener@tensorflow.org>: [GML] Implement dim-based shape reification for concat -- 467273958 by A. Unique TensorFlower<gardener@tensorflow.org>: Change mutexes under stream_executor/gpu to use absl::Mutex and absl::MutexLock instead of tensorflow::mutex and tensorflow::mutex_lock. Change instance of absl::make_unique to std::make_unique -- 467272897 by A. Unique TensorFlower<gardener@tensorflow.org>: [tf.data] Prepend `/bufferedio/` for all paths passed to LoadDataset op. -- PiperOrigin-RevId: 467408627
tensorflow · Aug 13, 2022 · 49f97f1 · 49f97f1
1 parent db84fb3
commit 49f97f1
Show file tree

Hide file tree

Showing 117 changed files with 3,908 additions and 1,137 deletions.
diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD
@@ -381,8 +381,10 @@ cc_library(
         "//tensorflow/python:__pkg__",
     ],
     deps = [
+        ":constants",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/platform:types",
+        "//tensorflow/core/util/tensor_bundle:naming",
         "@com_google_protobuf//:protobuf_headers",
         "//tensorflow/core/grappler:op_types",
         "@com_google_absl//absl/container:btree",

diff --git a/tensorflow/cc/saved_model/fingerprinting.cc b/tensorflow/cc/saved_model/fingerprinting.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/strip.h"
+#include "tensorflow/cc/saved_model/constants.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/op_def.pb.h"
@@ -31,13 +32,16 @@ limitations under the License.
 #include "tensorflow/core/grappler/op_types.h"
 #include "tensorflow/core/lib/strings/numbers.h"
 #include "tensorflow/core/lib/strings/proto_serialization.h"
+#include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/fingerprint.h"
+#include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/statusor.h"
 #include "tensorflow/core/protobuf/fingerprint.pb.h"
 #include "tensorflow/core/protobuf/meta_graph.pb.h"
 #include "tensorflow/core/protobuf/saved_model.pb.h"
 #include "tensorflow/core/protobuf/saved_object_graph.pb.h"
+#include "tensorflow/core/util/tensor_bundle/naming.h"
 
 namespace tensorflow::fingerprinting {
 
@@ -86,6 +90,20 @@ void CanonicalizeNodes(GraphDef* graph_def) {
   }
 }
 
+// Returns the hash of the checkpoint .index file, 0 if there is none.
+uint64 HashCheckpointIndexFile(absl::string_view model_dir) {
+  std::string meta_filename = MetaFilename(io::JoinPath(
+      model_dir, kSavedModelVariablesDirectory, kSavedModelVariablesFilename));
+  std::string data;
+  Status read_status = ReadFileToString(Env::Default(), meta_filename, &data);
+  if (read_status.ok()) {
+    return tensorflow::Fingerprint64(data);
+  } else {
+    LOG(WARNING) << read_status.error_message();
+    return 0;
+  }
+}
+
 }  // namespace
 
 uint64 ComputeHash(const GraphDef& graph_def) {
@@ -94,7 +112,8 @@ uint64 ComputeHash(const GraphDef& graph_def) {
   return tensorflow::Fingerprint64(graph_def_string);
 }
 
-FingerprintDef CreateFingerprintDef(const MetaGraphDef& metagraph) {
+FingerprintDef CreateFingerprintDef(const MetaGraphDef& metagraph,
+                                    absl::string_view export_dir) {
   // Create a copy of `metagraph` which will be used and mutated for fingerprint
   // computation.
   MetaGraphDef metagraph_copy = metagraph;
@@ -114,6 +133,8 @@ FingerprintDef CreateFingerprintDef(const MetaGraphDef& metagraph) {
       RegularizeAndHashSavedObjectGraph(metagraph_copy.object_graph_def());
   fingerprint_def.set_saved_object_graph_hash(
       RegularizeAndHashSavedObjectGraph(metagraph_copy.object_graph_def()));
+  // Set fingerprint field #5.
+  fingerprint_def.set_checkpoint_hash(HashCheckpointIndexFile(export_dir));
   // Set version of the fingerprint.
   VersionDef* version = fingerprint_def.mutable_version();
   version->set_producer(kFingerprintProducer);

diff --git a/tensorflow/cc/saved_model/fingerprinting.h b/tensorflow/cc/saved_model/fingerprinting.h
@@ -37,8 +37,10 @@ uint64 RegularizeAndHashSignatureDefs(
 uint64 RegularizeAndHashSavedObjectGraph(
     const SavedObjectGraph& object_graph_def);
 
-// Creates a FingerprintDef proto from a MetaGraph.
-FingerprintDef CreateFingerprintDef(const MetaGraphDef& metagraph);
+// Creates a FingerprintDef proto from a MetaGraph and the checkpoint meta file
+// (.index) in `export_dir`.
+FingerprintDef CreateFingerprintDef(const MetaGraphDef& metagraph,
+                                    absl::string_view export_dir);
 
 // Canonicalizes the GraphDef in order to remove sources of non-determinism.
 void CanonicalizeGraphDef(GraphDef& graph_def);

diff --git a/tensorflow/cc/saved_model/fingerprinting_test.cc b/tensorflow/cc/saved_model/fingerprinting_test.cc
@@ -43,10 +43,11 @@ GraphDef CreateTestProto() {
   return graph_def;
 }
 
-StatusOr<SavedModel> ReadSavedModel(absl::string_view file_path) {
+StatusOr<SavedModel> ReadSavedModel(absl::string_view file_dir) {
+  std::string file_path = io::JoinPath(file_dir, "saved_model.pb");
   std::string serialized_saved_model;
-  auto status = ReadFileToString(Env::Default(), std::string(file_path),
-                                 &serialized_saved_model);
+  auto status =
+      ReadFileToString(Env::Default(), file_path, &serialized_saved_model);
   if (!status.ok()) {
     return status;
   }
@@ -63,25 +64,26 @@ TEST(FingerprintingTest, TestComputeHash) {
 TEST(FingerprintingTest, TestCreateFingerprint) {
   const std::string export_dir =
       io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
-                   "VarsAndArithmeticObjectGraph", "saved_model.pb");
-
+                   "VarsAndArithmeticObjectGraph");
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
                           ReadSavedModel(export_dir));
   FingerprintDef fingerprint_def =
-      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0), export_dir);
 
   EXPECT_GT(fingerprint_def.graph_def_checksum(), 0);
   EXPECT_EQ(fingerprint_def.graph_def_program_hash(), 10127142238652115842U);
   EXPECT_EQ(fingerprint_def.signature_def_hash(), 5693392539583495303);
   EXPECT_EQ(fingerprint_def.saved_object_graph_hash(), 3678101440349108924);
+  // TODO(b/242348400): The checkpoint hash is non-deterministic, so we cannot
+  // check its value here.
+  EXPECT_GT(fingerprint_def.checkpoint_hash(), 0);
 }
 
 // Test that canonicalization returns the same hash for two models saved by
 // calling `tf.saved_model.save` twice in a row in the same program.
 TEST(FingerprintingTest, TestCanonicalizeGraphDeforModelSavedTwice) {
-  const std::string export_dir =
-      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
-                   "bert1", "saved_model.pb");
+  const std::string export_dir = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
                           ReadSavedModel(export_dir));
 
@@ -90,9 +92,8 @@ TEST(FingerprintingTest, TestCanonicalizeGraphDeforModelSavedTwice) {
   CanonicalizeGraphDef(*graph_def);
   uint64 hash1 = ComputeHash(*graph_def);
 
-  const std::string export_dir2 =
-      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
-                   "bert2", "saved_model.pb");
+  const std::string export_dir2 = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert2");
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb2,
                           ReadSavedModel(export_dir2));
   const MetaGraphDef& metagraph2 = saved_model_pb2.meta_graphs(0);
@@ -106,22 +107,20 @@ TEST(FingerprintingTest, TestCanonicalizeGraphDeforModelSavedTwice) {
 // Compare the fingerprints of two models saved by calling
 // `tf.saved_model.save` twice in a row in the same program.
 TEST(FingerprintingTest, TestCompareFingerprintForTwoModelSavedTwice) {
-  const std::string export_dir =
-      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
-                   "bert1", "saved_model.pb");
+  const std::string export_dir = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
 
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
                           ReadSavedModel(export_dir));
   FingerprintDef fingerprint_def =
-      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0), export_dir);
 
-  const std::string export_dir2 =
-      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
-                   "bert2", "saved_model.pb");
+  const std::string export_dir2 = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert2");
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb2,
                           ReadSavedModel(export_dir2));
   FingerprintDef fingerprint_def2 =
-      CreateFingerprintDef(saved_model_pb2.meta_graphs(0));
+      CreateFingerprintDef(saved_model_pb2.meta_graphs(0), export_dir2);
 
   EXPECT_EQ(fingerprint_def.graph_def_program_hash(),
             fingerprint_def2.graph_def_program_hash());
@@ -132,30 +131,38 @@ TEST(FingerprintingTest, TestCompareFingerprintForTwoModelSavedTwice) {
 }
 
 TEST(FingerprintingTest, TestFingerprintComputationDoesNotMutateModel) {
-  const std::string export_dir =
-      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
-                   "bert1", "saved_model.pb");
+  const std::string export_dir = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
                           ReadSavedModel(export_dir));
   FingerprintDef fingerprint_def =
-      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0), export_dir);
   FingerprintDef fingerprint_def2 =
-      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0), export_dir);
 
   EXPECT_EQ(fingerprint_def.graph_def_checksum(),
             fingerprint_def2.graph_def_checksum());
 }
 
 TEST(FingerprintingTest, TestFingerprintHasVersion) {
-  const std::string export_dir =
-      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
-                   "bert1", "saved_model.pb");
+  const std::string export_dir = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
                           ReadSavedModel(export_dir));
   FingerprintDef fingerprint_def =
-      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0), export_dir);
   EXPECT_EQ(fingerprint_def.version().producer(), 0);
 }
 
+TEST(FingerprintingTest, TestHashCheckpointForModelWithNoVariables) {
+  const std::string export_dir = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
+  TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
+                          ReadSavedModel(export_dir));
+  FingerprintDef fingerprint_def =
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0), export_dir);
+  EXPECT_EQ(fingerprint_def.checkpoint_hash(), 0);
+}
+
 }  // namespace
 }  // namespace tensorflow::fingerprinting
diff --git a/tensorflow/compiler/mlir/g3doc/_includes/tf_passes.md b/tensorflow/compiler/mlir/g3doc/_includes/tf_passes.md
@@ -541,6 +541,21 @@ will be transformed into this region-based operation
 This pass is performing fusion specific to GPU targets. This is an ad-hoc
 pass for now, but should be integrated with some notion of "target" in the
 MLIR pipeline in the future.
+### `-tf-group-by-dialect`: Groups ops into functions that only contain one dialect.
+Factors operations into subroutines such that all functions only
+contain a single dialect. Which of the dialects are allowed in the
+"top" function is configurable.
+
+For example, the code
+  x.a()
+  x.b()
+  %c = y.c()
+  x.d(%c)
+would be transformed into something like
+  call @x_1()
+  %c = call @y_1()
+  call @x_2(%c)
+with @x_1, @x_2 and @y_1 filled in.
 ### `-tf-guarantee-all-funcs-one-use`: Guarantee all FuncOp's have only a single use.
 ### `-tf-hoist-replicate-invariant-resource-writes`: Hoists writes to replicate invariant resource variables.
 This pass hoists replicate invariant resource variable writes outside
@@ -661,6 +676,20 @@ Would be transformed to:
 -direction             : Move transposes to the beginning or the end of the block where they are defined.
 ```
 ### `-tf-optimize`: Optimize TensorFlow module
+### `-tf-order-by-dialect`: Reorders ops so ops of the same dialect are next to each other.
+Performs a reordering of ops so that
+  (a) ops of the same dialect are next to each other
+  (b) order within a dialect is preserved
+.
+For example, this would transform
+  %a = "x.f"()
+  %b = "y.f"(%a)
+  %c = "x.f"(%a)
+to
+  %a = "x.f"()
+  %c = "x.f"(%a)
+  %b = "y.f"(%a)
+so that the two "x" dialect instructions are next to each other.
 ### `-tf-outside-compiled-to-host-launch`: Wraps each op with the _xla_outside_compiled attribute in a separate tf_device.launch on replicated host device.
 This pass wraps ops with the same `_xla_outside_compilation`
 attribute value in a tf_device.launch op with host device assignment.

diff --git a/tensorflow/compiler/mlir/lite/tf_tfl_translate.cc b/tensorflow/compiler/mlir/lite/tf_tfl_translate.cc
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "absl/strings/str_split.h"
 #include "llvm/ADT/None.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
@@ -294,10 +295,11 @@ int main(int argc, char **argv) {
   });
 
   std::string result;
-  // TODO(b/153507667): Pass the session object when importing logic is removed.
+  llvm::Optional<tensorflow::Session *> session = llvm::None;
+  if (bundle) session = bundle->GetSession();
   auto status = tensorflow::ConvertTFExecutorToTFLOrFlatbuffer(
       module.ValueOrDie().get(), output_mlir, toco_flags, pass_config, tags,
-      /*saved_model_dir=*/"", /*session=*/llvm::None, &result);
+      /*saved_model_dir=*/"", session, &result);
   if (!status.ok()) return kTrFailure;
 
   std::string error_msg;

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -1153,6 +1153,7 @@ cc_library(
         "transforms/generated_optimize.inc",
         "transforms/gpu_fusion.cc",
         "transforms/graph_pruning.cc",
+        "transforms/group_by_dialect.cc",
         "transforms/guarantee_all_funcs_one_use.cc",
         "transforms/hoist_replicate_invariant_resource_writes.cc",
         "transforms/host_launch_to_outside_compiled.cc",
@@ -1165,6 +1166,7 @@ cc_library(
         "transforms/materialize_mlir_passthrough_op.cc",
         "transforms/merge_control_flow.cc",
         "transforms/optimize.cc",
+        "transforms/order_by_dialect.cc",
         "transforms/outside_compiled_to_host_launch.cc",
         "transforms/parallel_execute_to_islands.cc",
         "transforms/prepare_tpu_computation_for_tf_export.cc",