chore: import upstream snapshot with attribution

tensorflow/compiler/jit/tests/BUILD

@@ -0,0 +1,126 @@
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+licenses(["notice"])
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:private"],
+)
+
+cc_library(
+    name = "auto_clustering_test_helper",
+    testonly = True,
+    srcs = ["auto_clustering_test_helper.cc"],
+    hdrs = ["auto_clustering_test_helper.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/jit:compilation_passes",
+        "//tensorflow/compiler/jit:jit_compilation_passes",
+        "//tensorflow/compiler/jit:xla_cluster_util",
+        "//tensorflow/compiler/jit:xla_cpu_jit",
+        "//tensorflow/compiler/jit:xla_gpu_jit",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/tools/optimization:optimization_pass_runner_lib",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@xla//xla:status_macros",
+    ],
+)
+
+tf_cc_test(
+    name = "auto_clustering_test",
+    srcs = ["auto_clustering_test.cc"],
+    data = [
+        "keras_imagenet_main.golden_summary",
+        "keras_imagenet_main.pbtxt",
+        "keras_imagenet_main_graph_mode.golden_summary",
+        "keras_imagenet_main_graph_mode.pbtxt",
+        "opens2s_gnmt_mixed_precision.golden_summary",
+        "opens2s_gnmt_mixed_precision.pbtxt.gz",
+    ],
+    tags = ["config-cuda-only"],
+    deps = [
+        ":auto_clustering_test_helper",
+        "//tensorflow/core:test",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "device_compiler_test_helper",
+    testonly = True,
+    srcs = ["device_compiler_test_helper.cc"],
+    hdrs = ["device_compiler_test_helper.h"],
+    visibility = [
+        "//tensorflow/compiler/jit:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/compiler/jit:xla_activity_listener",
+        "//tensorflow/compiler/jit:xla_compilation_cache_proto_cc",
+        "//tensorflow/compiler/jit:xla_cpu_jit",
+        "//tensorflow/compiler/jit:xla_gpu_device",
+        "//tensorflow/compiler/jit:xla_gpu_jit",
+        "//tensorflow/core:all_kernels",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:direct_session",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:ops",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:path",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@xla//xla/service:hlo_proto_cc",
+    ],
+)
+
+tf_cc_test(
+    name = "device_compiler_serialize_test",
+    srcs = [
+        "device_compiler_serialize_test.cc",
+    ],
+    tags = [
+        "config-cuda-only",
+        "no_oss",  # This test only runs with GPU.
+        "requires-gpu-nvidia",
+        "xla",
+    ],
+    deps = [
+        ":device_compiler_test_helper",
+        "//tensorflow/compiler/jit:compilation_passes",
+        "//tensorflow/compiler/jit:flags",
+        "//tensorflow/core:test",
+    ],
+)
+
+tf_cc_test(
+    name = "device_compiler_serialize_options_test",
+    srcs = [
+        "device_compiler_serialize_options_test.cc",
+    ],
+    tags = [
+        "config-cuda-only",
+        "no_oss",  # This test only runs with GPU.
+        "requires-gpu-nvidia",
+        "xla",
+    ],
+    deps = [
+        ":device_compiler_test_helper",
+        "//tensorflow/compiler/jit:compilation_passes",
+        "//tensorflow/compiler/jit:flags",
+        "//tensorflow/core:test",
+        "//tensorflow/core/framework:graph_proto_cc",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest",
+    ],
+)

tensorflow/compiler/jit/tests/auto_clustering_test.cc

@@ -0,0 +1,101 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/jit/tests/auto_clustering_test_helper.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+class AutoClusteringTestImpl : public AutoClusteringTest {
+ protected:
+  // Test auto-clustering with a proto text file ${key}.pbtxt.
+  absl::Status RunAutoClusteringTestWithPbtxt(absl::string_view key) {
+    std::string file_name_without_extension =
+        absl::StrCat(testing::TensorFlowSrcRoot(), "/compiler/jit/tests/", key);
+
+    return AutoClusteringTest::RunAutoClusteringTestWithPbtxt(
+        absl::StrCat(file_name_without_extension, ".pbtxt"),
+        absl::StrCat(file_name_without_extension, ".golden_summary"));
+  }
+
+  // Test auto-clustering with a gzipped proto text file ${key}.pbtxt.gz.
+  absl::Status RunAutoClusteringTestWithGzippedPbtxt(absl::string_view key) {
+    std::string file_name_without_extension =
+        absl::StrCat(testing::TensorFlowSrcRoot(), "/compiler/jit/tests/", key);
+
+    return AutoClusteringTest::RunAutoClusteringTestWithGzippedPbtxt(
+        absl::StrCat(file_name_without_extension, ".pbtxt.gz"),
+        absl::StrCat(file_name_without_extension, ".golden_summary"));
+  }
+};
+
+TEST_F(AutoClusteringTestImpl, KerasImagenetMain) {
+  // Generated from
+  //
+  //  TARGET_PATH=tensorflow_models/official/legacy/image_classification       \
+  //  bazel run -c opt --config=cuda ${TARGET_PATH}:resnet_imagenet_main       \
+  //    -- --skip_eval --num_gpus=1 --dtype=fp16 --batch_size=192              \
+  //    --train_steps=210 --enable_xla --enable_eager=true
+  //
+  // At CL 245846452
+  TF_ASSERT_OK(RunAutoClusteringTestWithPbtxt("keras_imagenet_main"));
+}
+
+TEST_F(AutoClusteringTestImpl, KerasImagenetMainGraphMode) {
+  // Generated from
+  //
+  //  TARGET_PATH=tensorflow_models/official/legacy/image_classification       \
+  //  bazel run -c opt --config=cuda ${TARGET_PATH}:resnet_imagenet_main       \
+  //   -- --use_synthetic_data --num_gpus=1 --batch_size=117 --train_steps=600 \
+  //   --skip_eval=True --logtostderr --enable_xla
+  TF_ASSERT_OK(
+      RunAutoClusteringTestWithPbtxt("keras_imagenet_main_graph_mode"));
+}
+
+TEST_F(AutoClusteringTestImpl, OpenSeq2SeqGNMT) {
+  // Model is from https://github.com/NVIDIA/OpenSeq2Seq.
+  // Generated from
+  //
+  // python run.py \
+  // --config_file=example_configs/text2text/en-de/en-de-gnmt-like-4GPUs.py \
+  // --use_xla_jit
+  TF_ASSERT_OK(
+      RunAutoClusteringTestWithGzippedPbtxt("opens2s_gnmt_mixed_precision"));
+}
+
+#if defined(PLATFORM_GOOGLE)
+absl::Status BenchmarkHelper(absl::string_view key, benchmark::State& state) {
+  return BenchmarkMarkForCompilation(
+      absl::StrCat(testing::TensorFlowSrcRoot(), "/compiler/jit/tests/", key,
+                   ".pbtxt"),
+      state);
+}
+
+void BM_MarkForCompilationPass_KerasImagenetMain(benchmark::State& state) {
+  CHECK_OK(BenchmarkHelper("keras_imagenet_main", state));
+}
+
+BENCHMARK(BM_MarkForCompilationPass_KerasImagenetMain);
+#endif  // PLATFORM_GOOGLE
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/jit/tests/auto_clustering_test_helper.cc

@@ -0,0 +1,234 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/tests/auto_clustering_test_helper.h"
+
+#include <memory>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/numbers.h"
+#include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
+#include "tensorflow/compiler/jit/xla_cluster_util.h"
+#include "xla/status_macros.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/io/random_inputstream.h"
+#include "tensorflow/core/lib/io/zlib_compression_options.h"
+#include "tensorflow/core/lib/io/zlib_inputstream.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+#include "tensorflow/core/util/port.h"
+#include "tensorflow/tools/optimization/optimization_pass_runner.h"
+
+namespace tensorflow {
+namespace {
+absl::StatusOr<std::string> SummarizeClustering(
+    const GraphDef& auto_clustered_graph_def) {
+  testing::ResetClusterSequenceNumber();
+  Graph graph(OpRegistry::Global());
+  GraphConstructorOptions graph_opts;
+  graph_opts.expect_device_spec = true;
+  graph_opts.allow_internal_ops = true;
+  TF_RETURN_IF_ERROR(
+      ConvertGraphDefToGraph(graph_opts, auto_clustered_graph_def, &graph));
+
+  // cluster_id -> (operation name -> # of operations)
+  const int kNoCluster = -1;
+  std::map<int, std::map<std::string, int>> clusters;
+  std::map<int, int> cluster_size;
+  int clustered_nodes = 0;
+  for (Node* n : graph.op_nodes()) {
+    int cluster = kNoCluster;
+    if (std::optional<absl::string_view> maybe_cluster =
+            GetXlaClusterForNode(*n)) {
+      maybe_cluster->remove_prefix(absl::string_view("cluster_").size());
+      TF_RET_CHECK(absl::SimpleAtoi(*maybe_cluster, &cluster));
+      clustered_nodes++;
+    }
+    clusters[cluster][n->type_string()]++;
+    cluster_size[cluster]++;
+  }
+
+  std::string result =
+      absl::StrCat("Clustered nodes: ", clustered_nodes,
+                   "\nUnclustered nodes: ", cluster_size[kNoCluster],
+                   "\nNumber of clusters: ", clusters.size() - 1, "\n\n");
+  for (const auto& pair : clusters) {
+    if (pair.first == kNoCluster) {
+      absl::StrAppend(&result, "unclustered");
+    } else {
+      absl::StrAppend(&result, "cluster ", pair.first);
+    }
+
+    absl::StrAppend(&result, " size ", cluster_size[pair.first], "\n");
+
+    for (const auto& ops_and_counts : pair.second) {
+      absl::StrAppend(&result, " ", ops_and_counts.first, " ",
+                      ops_and_counts.second, "\n");
+    }
+  }
+
+  return result;
+}
+
+absl::Status AssertGraphDefIsUnclustered(const GraphDef& graphdef) {
+  const char* kXlaClusterAttr = "_XlaCluster";
+  const char* kXlaAlreadyClusteredAttr = "_XlaAlreadyClustered";
+
+  for (const NodeDef& node : graphdef.node()) {
+    if (node.attr().count(kXlaClusterAttr) ||
+        node.attr().count(kXlaAlreadyClusteredAttr)) {
+      return absl::InvalidArgumentError(
+          "Input files are already clustered, you probably copied in "
+          "mark_for_compilation_<n>.pbtxt when you should have copied in "
+          "before_mark_for_compilation_<n>.pbtxt");
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status ReadTextProtoFromString(Env* env, const std::string& data,
+                                     ::tensorflow::protobuf::Message* proto) {
+  if (!::tensorflow::protobuf::TextFormat::ParseFromString(data, proto)) {
+    return absl::DataLossError("Can't parse input data as text proto");
+  }
+  return absl::OkStatus();
+}
+}  // namespace
+
+absl::Status AutoClusteringTest::RunAutoClusteringTestImpl(
+    GraphDef graphdef, absl::string_view golden_summary_file_path) {
+  if (!IsGoogleCudaEnabled()) {
+    // There is some slight change in the clustering decisions under
+    // --config=cuda.  I have not looked closely at why that is happening, but
+    // most likely some of the partial declustering passes behave differently
+    // with --config=cuda because of different HostMemory.  So for now only test
+    // the non-CUDA config, under the assumption that regressions with
+    // --config=cuda would also be detected as regressions without
+    // --config=cuda.
+
+    LOG(INFO) << "Not running "
+              << ::testing::UnitTest::GetInstance()->current_test_info()->name()
+              << " since test was not built with --config=cuda";
+    return absl::OkStatus();
+  }
+
+  TF_RETURN_IF_ERROR(AssertGraphDefIsUnclustered(graphdef));
+
+  OptimizationPassRunner runner;
+  TF_RETURN_IF_ERROR(runner.SetJitLevel(tensorflow::OptimizerOptions::ON_2));
+  TF_RETURN_IF_ERROR(runner.AddCpus(32));
+  TF_RETURN_IF_ERROR(runner.AddGpus(8));
+
+  for (absl::string_view auto_clustering_pass :
+       {"CloneConstantsForBetterClusteringPass", "MarkForCompilationPass",
+        "IncreaseDynamismForAutoJitPass", "PartiallyDeclusterPass"}) {
+    GraphDef next;
+    TF_RETURN_IF_ERROR(
+        runner.Run(auto_clustering_pass, std::move(graphdef), &next));
+    graphdef = std::move(next);
+  }
+
+  TF_ASSIGN_OR_RETURN(std::string clustering_summary,
+                      SummarizeClustering(graphdef));
+
+  // To update golden files flip this to true and run
+  //
+  // bazel test --test_strategy=local \
+  //   tensorflow/compiler/jit/tests:auto_clustering_test
+  bool update_golden = false;
+  if (update_golden) {
+    TF_RETURN_IF_ERROR(WriteStringToFile(Env::Default(),
+                                         std::string(golden_summary_file_path),
+                                         clustering_summary));
+  }
+
+  std::string golden_file_contents;
+  TF_RETURN_IF_ERROR(ReadFileToString(Env::Default(),
+                                      std::string(golden_summary_file_path),
+                                      &golden_file_contents));
+
+  EXPECT_EQ(golden_file_contents, clustering_summary);
+
+  return absl::OkStatus();
+}
+
+absl::Status AutoClusteringTest::RunAutoClusteringTestWithPbtxt(
+    absl::string_view pbtxt_file_path,
+    absl::string_view golden_summary_file_path) {
+  GraphDef graphdef;
+  TF_RETURN_IF_ERROR(
+      ReadTextProto(Env::Default(), std::string(pbtxt_file_path), &graphdef));
+  return RunAutoClusteringTestImpl(std::move(graphdef),
+                                   golden_summary_file_path);
+}
+
+absl::Status AutoClusteringTest::RunAutoClusteringTestWithGzippedPbtxt(
+    absl::string_view gzipped_pbtxt_file_path,
+    absl::string_view golden_summary_file_path) {
+  Env* env = Env::Default();
+  std::unique_ptr<RandomAccessFile> file_reader;
+  TF_RETURN_IF_ERROR(env->NewRandomAccessFile(
+      std::string(gzipped_pbtxt_file_path), &file_reader));
+  std::unique_ptr<io::RandomAccessInputStream> input_stream =
+      std::make_unique<io::RandomAccessInputStream>(file_reader.get());
+  constexpr int k_buffer_size = 256 << 10;  // 256kb
+  io::ZlibInputStream in(input_stream.get(),
+                         /*input_buffer_bytes=*/k_buffer_size,
+                         /*output_buffer_bytes=*/k_buffer_size,
+                         io::ZlibCompressionOptions::GZIP());
+  tstring decompressed_pbtxt_string;
+  absl::Status s = in.ReadNBytes(INT_MAX, &decompressed_pbtxt_string);
+  if (!s.ok() && !absl::IsOutOfRange(s)) {
+    // OutOfRange is fine since we set the number of read bytes to INT_MAX.
+    // Only return other kinds of errors.
+    return s;
+  }
+
+  GraphDef graphdef;
+  TF_RETURN_IF_ERROR(ReadTextProtoFromString(
+      Env::Default(), decompressed_pbtxt_string, &graphdef));
+  return RunAutoClusteringTestImpl(std::move(graphdef),
+                                   golden_summary_file_path);
+}
+
+#if defined(PLATFORM_GOOGLE)
+absl::Status BenchmarkMarkForCompilation(absl::string_view graph_def_path,
+                                         benchmark::State& state) {
+  GraphDef graph_def;
+  TF_RETURN_IF_ERROR(
+      ReadTextProto(Env::Default(), std::string(graph_def_path), &graph_def));
+
+  OptimizationPassRunner runner;
+  TF_RETURN_IF_ERROR(runner.SetJitLevel(tensorflow::OptimizerOptions::ON_2));
+  TF_RETURN_IF_ERROR(runner.AddCpus(32));
+  TF_RETURN_IF_ERROR(runner.AddGpus(8));
+
+  for (auto _ : state) {
+    state.PauseTiming();
+    GraphDef result;
+    GraphDef graph_def_copy = graph_def;
+    state.ResumeTiming();
+    TF_RETURN_IF_ERROR(runner.Run("MarkForCompilationPass",
+                                  std::move(graph_def_copy), &result));
+  }
+
+  return absl::OkStatus();
+}
+#endif  // PLATFORM_GOOGLE
+
+}  // namespace tensorflow

tensorflow/compiler/jit/tests/auto_clustering_test_helper.h

@@ -0,0 +1,71 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_JIT_TESTS_AUTO_CLUSTERING_TEST_HELPER_H_
+#define TENSORFLOW_COMPILER_JIT_TESTS_AUTO_CLUSTERING_TEST_HELPER_H_
+
+#include "absl/status/statusor.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace tensorflow {
+// Helper to write integration tests and benchmarks for the auto-clustering pass
+// pipeline.  These tests run auto-clustering on a graphdef and compare a
+// summary of the auto-clustering decisions with a "golden" summary.
+//
+// To create a new test from an TF workload first run the workload with the
+// following environment variables set:
+//
+//   TF_DUMP_GRAPH_PREFIX=<some temporary directory>
+//   TF_XLA_FLAGS="--tf_xla_clustering_debug"
+//
+// If auto-clustering is enabled this should produce files named
+// before_mark_for_compilation_<N>.pbtxt in the temporary directory.  As the
+// file name suggests, these are graphdefs that have been dumped right before
+// the mark_for_compilation pass.  There should be one
+// before_mark_for_compilation_<N>.pbtxt for every TF graph that was
+// auto-clustered, out of which usually only one is the "main" graph that's
+// running training/inference.
+//
+// Copy the pbtxt for that "main" graph to tensorflow/compiler/jit/tests/
+// (i.e. this directory) and create a corresponding empty .golden_summary file.
+// Add the .pbtxt and .golden_summary files to the "data" section of the cc_test
+// rule for :auto_clustering_test and then see the comment on update_golden on
+// how to auto-generate the .golden_summary file.
+
+class AutoClusteringTest : public ::testing::Test {
+ protected:
+  absl::Status RunAutoClusteringTestWithPbtxt(
+      absl::string_view pbtxt_file_path,
+      absl::string_view golden_summary_file_path);
+  absl::Status RunAutoClusteringTestWithGzippedPbtxt(
+      absl::string_view gzipped_pbtxt_file_path,
+      absl::string_view golden_summary_file_path);
+
+ private:
+  absl::Status RunAutoClusteringTestImpl(
+      GraphDef graphdef, absl::string_view golden_summary_file_path);
+};
+
+#if defined(PLATFORM_GOOGLE)
+// Reads the GraphDef stored in graph_def_path (which must be a pbtxt file) and
+// benchmarks MarkForCompilationPass on this graphdef.
+absl::Status BenchmarkMarkForCompilation(absl::string_view graph_def_path,
+                                         benchmark::State& state);
+#endif  // PLATFORM_GOOGLE
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_JIT_TESTS_AUTO_CLUSTERING_TEST_HELPER_H_

tensorflow/compiler/jit/tests/device_compiler_serialize_options_test.cc

@@ -0,0 +1,78 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <gtest/gtest.h>
+#include "absl/strings/match.h"
+#include "tensorflow/compiler/jit/flags.h"
+#include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
+#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+TEST_F(DeviceCompilerSerializeTest, PersistentCacheOptionsTest) {
+  GraphDef graph = GetTestGraph({-1, 4});
+
+  // Warmup the persistent cache(s) with multiple runs. 4 is a magic number to
+  // detect non-determinism in TF when running the test.
+  listener()->ClearListenerHistory();
+  for (int b = 1; b < 4; ++b) {
+    TF_ASSERT_OK(ExecuteWithBatch(graph, b));
+  }
+  TF_ASSERT_OK(listener()->VerifyPersistentCacheUseListenerHistory(
+      /*expect_persistent_cache_use=*/false));
+
+  // Reset the cluster numbering between sessions so we can get the same
+  // cluster numbering.
+  testing::ResetClusterSequenceNumber();
+
+  auto status =
+      AlterPersistentCacheEntryHloModuleNames(tensorflow::testing::TmpDir());
+  EXPECT_FALSE(status.ok());
+  EXPECT_TRUE(absl::StrContains(
+      status.message(),
+      "Did not find any persistent XLA compilation cache entries to alter."));
+
+  TF_ASSERT_OK(AlterPersistentCacheEntryHloModuleNames(
+      tensorflow::testing::TmpDir(), "my_test_prefix"));
+
+  // Run again and these should all hit in the persistent cache despite having
+  // altered the persistent cache entries' HLO modules (disabled strict
+  // signature checks).
+  listener()->ClearListenerHistory();
+  for (int b = 1; b < 4; ++b) {
+    TF_ASSERT_OK(ExecuteWithBatch(graph, b));
+  }
+  TF_ASSERT_OK(listener()->VerifyPersistentCacheUseListenerHistory(
+      /*expect_persistent_cache_use=*/true));
+}
+
+}  // namespace
+}  // namespace tensorflow
+
+int main(int argc, char** argv) {
+  tensorflow::GetMarkForCompilationPassFlags()
+      ->tf_xla_deterministic_cluster_names = true;
+  tensorflow::GetMarkForCompilationPassFlags()
+      ->tf_xla_persistent_cache_directory = tensorflow::testing::TmpDir();
+  tensorflow::GetMarkForCompilationPassFlags()
+      ->tf_xla_disable_strict_signature_checks = true;
+  tensorflow::GetMarkForCompilationPassFlags()->tf_xla_persistent_cache_prefix =
+      "my_test_prefix";
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

tensorflow/compiler/jit/tests/device_compiler_serialize_test.cc

@@ -0,0 +1,75 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/flags.h"
+#include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
+#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+TEST_F(DeviceCompilerSerializeTest, PersistentCacheTest) {
+  GraphDef graph = GetTestGraph({-1, 4});
+
+  // Warmup the persistent cache(s) with multiple runs. 4 is a magic number to
+  // detect non-determinism in TF when running the test.
+  listener()->ClearListenerHistory();
+  for (int b = 1; b < 4; ++b) {
+    TF_ASSERT_OK(ExecuteWithBatch(graph, b));
+  }
+  TF_ASSERT_OK(listener()->VerifyPersistentCacheUseListenerHistory(
+      /*expect_persistent_cache_use=*/false));
+
+  // Reset the cluster numbering between sessions so we can get the same
+  // cluster numbering.
+  testing::ResetClusterSequenceNumber();
+
+  // Run again but these should all hit in the persistent cache.
+  listener()->ClearListenerHistory();
+  for (int b = 1; b < 4; ++b) {
+    TF_ASSERT_OK(ExecuteWithBatch(graph, b));
+  }
+  TF_ASSERT_OK(listener()->VerifyPersistentCacheUseListenerHistory(
+      /*expect_persistent_cache_use=*/true));
+
+  // Reset the cluster numbering between sessions so we can get the same
+  // cluster numbering.
+  testing::ResetClusterSequenceNumber();
+
+  TF_ASSERT_OK(
+      AlterPersistentCacheEntryHloModuleNames(tensorflow::testing::TmpDir()));
+
+  // Run again but these should all fail, because the persistent cache entries'
+  // HLO modules have been altered.
+  for (int b = 1; b < 4; ++b) {
+    auto status = ExecuteWithBatch(graph, b);
+    EXPECT_FALSE(status.ok());
+    EXPECT_TRUE(
+        absl::StrContains(status.message(), "Serialized HLO does not match."));
+  }
+}
+
+}  // namespace
+}  // namespace tensorflow
+
+int main(int argc, char** argv) {
+  tensorflow::GetMarkForCompilationPassFlags()
+      ->tf_xla_deterministic_cluster_names = true;
+  tensorflow::GetMarkForCompilationPassFlags()
+      ->tf_xla_persistent_cache_directory = tensorflow::testing::TmpDir();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

tensorflow/compiler/jit/tests/device_compiler_test_helper.cc

@@ -0,0 +1,167 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"
+
+#include <string>
+
+#include "absl/strings/match.h"
+#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
+#include "xla/service/hlo.pb.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/public/session.h"
+
+namespace tensorflow {
+namespace {
+
+// Creates a float tensor of linearly increasing values, starting from offset.
+Tensor CreateInputTensor(const TensorShape& shape, float offset) {
+  Tensor tensor(DT_FLOAT, shape);
+  for (int64_t i = 0; i < tensor.flat<float>().size(); ++i) {
+    tensor.flat<float>()(i) = offset + i;
+  }
+  return tensor;
+}
+
+NodeDef MakeNode(
+    absl::string_view name, absl::string_view op,
+    absl::Span<const std::string> inputs,
+    absl::Span<
+        const std::pair<std::string, FunctionDefHelper::AttrValueWrapper>>
+        attrs) {
+  NodeDef node;
+  node.set_name(std::string(name));
+  node.set_op(std::string(op));
+  for (const auto& input : inputs) node.add_input(input);
+  for (const auto& attr : attrs)
+    node.mutable_attr()->insert({attr.first, attr.second.proto});
+  return node;
+}
+
+}  // namespace
+
+GraphDef DeviceCompilerSerializeTest::GetTestGraph(
+    const PartialTensorShape& input_shape) {
+  FunctionDef make_test_fn = FunctionDefHelper::Define(
+      "TestFn", {"a:float", "b:float", "c:float"}, {"m:float"}, {},
+      {{{"d"}, "Add", {"a", "b"}, {{"T", DT_FLOAT}}},
+       {{"e"}, "Mul", {"d", "c"}, {{"T", DT_FLOAT}}},
+       {{"f"}, "Add", {"e", "a"}, {{"T", DT_FLOAT}}},
+       {{"g"}, "Mul", {"f", "b"}, {{"T", DT_FLOAT}}},
+       // Force two clusters by excluding this node explicitly.
+       {{"h"}, "Add", {"g", "f"}, {{"T", DT_FLOAT}, {"_XlaCompile", false}}},
+       {{"i"}, "Add", {"h", "e"}, {{"T", DT_FLOAT}}},
+       {{"j"}, "Add", {"i", "h"}, {{"T", DT_FLOAT}}},
+       {{"k"}, "Add", {"j", "h"}, {{"T", DT_FLOAT}}},
+       {{"l"}, "Add", {"k", "h"}, {{"T", DT_FLOAT}}},
+       {{"m"}, "Identity", {"l"}, {{"T", DT_FLOAT}}}});
+
+  GraphDef graph;
+  *graph.mutable_library()->add_function() = make_test_fn;
+  *graph.add_node() = MakeNode("a", "Placeholder", {},
+                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
+  *graph.add_node() = MakeNode("b", "Placeholder", {},
+                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
+  *graph.add_node() = MakeNode("c", "Placeholder", {},
+                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
+  *graph.add_node() = MakeNode("m", "TestFn", {"a", "b", "c"}, {});
+  return graph;
+}
+
+absl::Status DeviceCompilerSerializeTest::ExecuteWithBatch(
+    const GraphDef& graph, int batch) {
+  const TensorShape shape({batch, 4});
+
+  // Compute the golden output tensor
+  std::vector<Tensor> golden_output_tensors;
+  {
+    SessionOptions options;
+    std::unique_ptr<Session> session(NewSession(options));
+    TF_RETURN_IF_ERROR(session->Create(graph));
+    RunOptions run_options;
+
+    Tensor input_a = CreateInputTensor(shape, 0);
+    Tensor input_b = CreateInputTensor(shape, shape.num_elements());
+    Tensor input_c = CreateInputTensor(shape, 2 * shape.num_elements());
+    TF_RETURN_IF_ERROR(session->Run(
+        run_options,
+        {std::make_pair("a", input_a), std::make_pair("b", input_b),
+         std::make_pair("c", input_c)},
+        {"m"}, {}, &golden_output_tensors, nullptr));
+    TF_RETURN_IF_ERROR(session->Close());
+  }
+
+  // Compute the XLA compiled output
+  std::vector<Tensor> output_tensors;
+  {
+    SessionOptions options;
+    auto& opts =
+        *options.config.mutable_graph_options()->mutable_optimizer_options();
+    opts.set_global_jit_level(OptimizerOptions::ON_1);
+    opts.set_cpu_global_jit(true);
+
+    std::unique_ptr<Session> session(NewSession(options));
+    TF_RETURN_IF_ERROR(session->Create(graph));
+    RunOptions run_options;
+    Tensor input_a = CreateInputTensor(shape, 0);
+    Tensor input_b = CreateInputTensor(shape, shape.num_elements());
+    Tensor input_c = CreateInputTensor(shape, 2 * shape.num_elements());
+    TF_RETURN_IF_ERROR(session->Run(
+        run_options,
+        {std::make_pair("a", input_a), std::make_pair("b", input_b),
+         std::make_pair("c", input_c)},
+        {"m"}, {}, &output_tensors, nullptr));
+    TF_RETURN_IF_ERROR(session->Close());
+  }
+
+  Tensor f32_input(DT_FLOAT, shape);
+  for (int64_t i = 0; i < f32_input.NumElements(); ++i) {
+    EXPECT_NEAR(golden_output_tensors[0].flat<float>()(i),
+                output_tensors[0].flat<float>()(i), 1e-3);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status
+DeviceCompilerSerializeTest::AlterPersistentCacheEntryHloModuleNames(
+    absl::string_view persistent_cache_dir_path,
+    absl::string_view file_prefix) {
+  Env* env = Env::Default();
+  std::vector<std::string> file_names;
+  TF_RETURN_IF_ERROR(
+      env->GetChildren(tensorflow::testing::TmpDir(), &file_names));
+
+  bool altered = false;
+  for (const auto& file_name : file_names) {
+    if (absl::EndsWith(file_name, ".pb") &&
+        absl::StartsWith(file_name, file_prefix)) {
+      XlaSerializedCacheEntry entry;
+      auto file_path = io::JoinPath(persistent_cache_dir_path, file_name);
+      TF_RETURN_IF_ERROR(ReadTextOrBinaryProto(env, file_path, &entry));
+      entry.mutable_hlo_module()->set_name(
+          absl::StrCat(entry.hlo_module().name(), "_altered"));
+      TF_RETURN_IF_ERROR(WriteBinaryProto(env, file_path, entry));
+      altered = true;
+    }
+  }
+
+  if (!altered) {
+    return absl::NotFoundError(
+        "Did not find any persistent XLA compilation cache entries to alter.");
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/jit/tests/device_compiler_test_helper.h

@@ -0,0 +1,104 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_JIT_TESTS_DEVICE_COMPILER_TEST_HELPER_H_
+#define TENSORFLOW_COMPILER_JIT_TESTS_DEVICE_COMPILER_TEST_HELPER_H_
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/jit/xla_activity_listener.h"
+#include "tensorflow/core/graph/graph_def_builder.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+// A listener to inspect the use of XLA's persistent compilation cache entries.
+class JitCompilationListener : public XlaActivityListener {
+ public:
+  absl::Status Listen(
+      const XlaAutoClusteringActivity& auto_clustering_activity) override {
+    return absl::OkStatus();
+  }
+
+  absl::Status Listen(
+      const XlaJitCompilationActivity& jit_compilation_activity) override {
+    activity_history_.push_back(jit_compilation_activity);
+    return absl::OkStatus();
+  }
+
+  absl::Status Listen(
+      const XlaOptimizationRemark& optimization_remark) override {
+    return absl::OkStatus();
+  }
+
+  ~JitCompilationListener() override = default;
+
+  absl::Status VerifyPersistentCacheUseListenerHistory(
+      bool expect_persistent_cache_use) {
+    for (const auto& activity : activity_history_) {
+      if (activity.used_persistent_cache() != expect_persistent_cache_use) {
+        return absl::FailedPreconditionError("Unexpected listener history.");
+      }
+    }
+    return absl::OkStatus();
+  }
+
+  std::vector<XlaJitCompilationActivity> GetListenerHistory() {
+    return activity_history_;
+  }
+
+  void ClearListenerHistory() { activity_history_.clear(); }
+
+ private:
+  std::vector<XlaJitCompilationActivity> activity_history_;
+};
+
+// Fixture for testing XLA compilation cache serialization.
+class DeviceCompilerSerializeTest : public ::testing::Test {
+ protected:
+  DeviceCompilerSerializeTest() {
+    auto listener = std::make_unique<JitCompilationListener>();
+    listener_ = listener.get();
+    RegisterXlaActivityListener(std::move(listener));
+  }
+
+  JitCompilationListener* listener() const { return listener_; }
+
+  // Returns a test graph that will split into two XLA clusters (due to a node
+  // with _XlaCompile = false).
+  GraphDef GetTestGraph(const PartialTensorShape& input_shape);
+
+  // Runs the graph using specified batch size both with and without XLA JIT
+  // compilation. Returns an error if the results between the two do not match.
+  absl::Status ExecuteWithBatch(const GraphDef& graph, int batch);
+
+  // Adds the suffix "_altered" to the HLO module names of all of the persistent
+  // XLA compilation cache entries found at the specified directory. If none are
+  // found, returns NOT_FOUND error.
+  absl::Status AlterPersistentCacheEntryHloModuleNames(
+      absl::string_view persistent_cache_dir_path,
+      absl::string_view file_prefix = "xla_compile_cache");
+
+ private:
+  JitCompilationListener* listener_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_JIT_TESTS_DEVICE_COMPILER_TEST_HELPER_H_

tensorflow/compiler/jit/tests/keras_imagenet_main.golden_summary

@@ -0,0 +1,60 @@
+Clustered nodes: 2725
+Unclustered nodes: 606
+Number of clusters: 2
+
+unclustered size 606
+ AssignAddVariableOp 1
+ Const 108
+ DivNoNan 1
+ Identity 2
+ Merge 53
+ PlaceholderWithDefault 1
+ ReadVariableOp 2
+ Switch 1
+ _Arg 435
+ _Retval 2
+cluster 0 size 1910
+ Add 16
+ AddN 71
+ ArgMax 1
+ AssignAddVariableOp 1
+ BiasAdd 1
+ BiasAddGrad 1
+ Cast 115
+ Const 407
+ Conv2D 53
+ Conv2DBackpropFilter 53
+ Conv2DBackpropInput 52
+ Equal 1
+ FusedBatchNormGradV2 53
+ FusedBatchNormV2 53
+ MatMul 3
+ MaxPool 1
+ MaxPoolGrad 1
+ Mean 1
+ Mul 218
+ Pad 2
+ ReadVariableOp 538
+ Relu 49
+ ReluGrad 49
+ Reshape 2
+ ResourceApplyKerasMomentum 161
+ Slice 1
+ Softmax 1
+ SparseSoftmaxCrossEntropyWithLogits 1
+ Squeeze 1
+ Sum 1
+ Tile 1
+ Transpose 1
+cluster 1 size 815
+ AddN 1
+ AssignAddVariableOp 1
+ AssignSubVariableOp 106
+ Const 220
+ DivNoNan 1
+ Identity 1
+ Mul 161
+ ReadVariableOp 106
+ Square 55
+ Sub 106
+ Sum 57

tensorflow/compiler/jit/tests/keras_imagenet_main_graph_mode.golden_summary

@@ -0,0 +1,51 @@
+Clustered nodes: 2178
+Unclustered nodes: 446
+Number of clusters: 1
+
+unclustered size 446
+ AssignAddVariableOp 2
+ Const 2
+ DivNoNan 1
+ Identity 1
+ NoOp 1
+ ReadVariableOp 2
+ VarHandleOp 435
+ _Retval 2
+cluster 0 size 2178
+ Add 17
+ AddN 72
+ ArgMax 1
+ AssignAddVariableOp 1
+ AssignSubVariableOp 106
+ BiasAdd 1
+ BiasAddGrad 1
+ Cast 3
+ Const 357
+ Conv2D 53
+ Conv2DBackpropFilter 53
+ Conv2DBackpropInput 52
+ DivNoNan 1
+ Equal 1
+ FusedBatchNorm 53
+ FusedBatchNormGrad 53
+ Identity 2
+ MatMul 3
+ MaxPool 1
+ MaxPoolGrad 1
+ Mean 1
+ Mul 164
+ Pad 1
+ ReadVariableOp 646
+ Relu 49
+ ReluGrad 49
+ Reshape 2
+ ResourceApplyKerasMomentum 161
+ ShapeN 50
+ Softmax 1
+ SparseSoftmaxCrossEntropyWithLogits 1
+ Square 55
+ Squeeze 1
+ Sub 106
+ Sum 57
+ Tile 1
+ Transpose 1

tensorflow/compiler/jit/tests/opens2s_gnmt_mixed_precision.golden_summary

@@ -0,0 +1,438 @@
+Clustered nodes: 2385
+Unclustered nodes: 4221
+Number of clusters: 30
+
+unclustered size 4221
+ Add 17
+ AddN 1
+ All 1
+ ApplyAdam 38
+ Assert 7
+ Assign 47
+ AssignAdd 2
+ AssignSub 2
+ BroadcastGradientArgs 44
+ Cast 38
+ ConcatV2 3
+ Const 875
+ ControlTrigger 5
+ Enter 874
+ Equal 4
+ Exit 69
+ ExpandDims 9
+ Fill 5
+ FloorMod 1
+ GreaterEqual 7
+ Identity 113
+ IsVariableInitialized 1
+ IteratorGetNext 1
+ IteratorV2 1
+ Less 9
+ LogicalAnd 3
+ LoopCond 8
+ Max 4
+ Maximum 44
+ Merge 145
+ Minimum 43
+ Mul 8
+ NextIteration 136
+ RandomUniform 14
+ Range 9
+ RefSwitch 166
+ Reshape 2
+ ScatterAdd 4
+ Shape 6
+ ShapeN 10
+ Size 2
+ Snapshot 1
+ StackPopV2 314
+ StackPushV2 314
+ StackV2 314
+ StridedSlice 12
+ Sub 5
+ Sum 2
+ Switch 247
+ TensorArrayGatherV3 8
+ TensorArrayGradV3 17
+ TensorArrayReadV3 9
+ TensorArrayScatterV3 8
+ TensorArraySizeV3 4
+ TensorArrayV3 8
+ TensorArrayWriteV3 9
+ Unique 2
+ VariableV2 164
+ _Retval 5
+cluster 0 size 440
+ Abs 40
+ AddN 1
+ Any 41
+ Cast 40
+ ConcatV2 2
+ Const 95
+ ExpandDims 2
+ IsInf 1
+ IsNan 40
+ L2Loss 40
+ LogicalOr 1
+ Max 41
+ Minimum 1
+ Mul 82
+ Pack 3
+ Reciprocal 2
+ Reshape 2
+ ReverseSequence 1
+ Sqrt 1
+ Sum 1
+ Transpose 3
+cluster 1 size 86
+ BroadcastGradientArgs 1
+ Cast 5
+ ConcatV2 1
+ Const 30
+ ExpandDims 3
+ Less 1
+ Mean 2
+ Minimum 1
+ Mul 3
+ Pack 2
+ Pad 1
+ Range 1
+ RealDiv 1
+ Reshape 8
+ Shape 7
+ Size 1
+ Slice 2
+ Snapshot 5
+ SparseSoftmaxCrossEntropyWithLogits 1
+ StridedSlice 2
+ Sub 3
+ Sum 1
+ Tile 1
+ Transpose 2
+ Unpack 1
+cluster 2 size 44
+ Cast 2
+ ConcatV2 2
+ Const 18
+ ExpandDims 1
+ GatherV2 2
+ Less 1
+ MatMul 1
+ Mul 1
+ Pack 1
+ Prod 2
+ Range 1
+ Reshape 3
+ Shape 8
+ StridedSlice 1
+cluster 3 size 10
+ AddN 1
+ Const 1
+ MatMul 2
+ Mul 1
+ Reshape 3
+ Sum 1
+ Transpose 1
+cluster 4 size 11
+ ConcatOffset 1
+ Const 4
+ ReverseSequence 1
+ Slice 2
+ Transpose 3
+cluster 5 size 21
+ All 1
+ ConcatV2 1
+ Const 11
+ Equal 1
+ ExpandDims 1
+ ReverseSequence 1
+ Shape 1
+ StridedSlice 1
+ Transpose 3
+cluster 6 size 11
+ Cast 1
+ Const 5
+ GatherV2 1
+ Shape 1
+ StridedSlice 1
+ Transpose 1
+ ZerosLike 1
+cluster 7 size 33
+ All 2
+ Cast 1
+ Const 17
+ Equal 2
+ ExpandDims 2
+ GatherV2 1
+ Identity 1
+ ReverseSequence 1
+ Shape 2
+ StridedSlice 2
+ Transpose 2
+cluster 8 size 11
+ Cast 1
+ Const 4
+ Floor 1
+ Identity 1
+ Mul 2
+ Pow 1
+ Sub 1
+cluster 9 size 5
+ All 1
+ Const 1
+ Less 1
+ LogicalAnd 1
+ LogicalNot 1
+cluster 10 size 9
+ All 1
+ Const 4
+ Equal 1
+ LessEqual 1
+ LogicalOr 1
+ Max 1
+cluster 11 size 302
+ Add 24
+ BatchMatMulV2 1
+ BiasAdd 8
+ Cast 8
+ ConcatV2 16
+ Const 81
+ ExpandDims 3
+ Fill 1
+ GreaterEqual 8
+ Identity 1
+ Less 1
+ MatMul 10
+ Mul 44
+ Range 1
+ Rsqrt 1
+ Select 19
+ Shape 6
+ Sigmoid 24
+ Snapshot 8
+ Softmax 1
+ Split 8
+ SplitV 6
+ Square 1
+ Squeeze 1
+ StridedSlice 1
+ Sum 2
+ Tanh 17
+cluster 12 size 6
+ Add 1
+ All 1
+ Const 2
+ GreaterEqual 1
+ LogicalOr 1
+cluster 15 size 614
+ Add 22
+ AddN 41
+ BatchMatMulV2 2
+ BiasAddGrad 8
+ BroadcastGradientArgs 4
+ ConcatOffset 1
+ ConcatV2 14
+ Const 199
+ DynamicStitch 1
+ FloorDiv 1
+ MatMul 20
+ Maximum 1
+ Mul 74
+ NoOp 13
+ Reshape 86
+ RsqrtGrad 1
+ Select 36
+ SigmoidGrad 24
+ Slice 32
+ StridedSlice 2
+ Sub 1
+ Sum 11
+ TanhGrad 17
+ Tile 2
+ ZerosLike 1
+cluster 16 size 22
+ Add 2
+ BiasAdd 1
+ ConcatV2 1
+ Const 3
+ GreaterEqual 1
+ MatMul 1
+ Mul 3
+ Select 3
+ Shape 1
+ Sigmoid 3
+ Split 1
+ Tanh 2
+cluster 17 size 60
+ Add 2
+ AddN 4
+ BiasAddGrad 1
+ BroadcastGradientArgs 1
+ Cast 2
+ ConcatOffset 1
+ ConcatV2 1
+ Const 8
+ MatMul 2
+ Mul 6
+ NoOp 2
+ Reshape 9
+ Select 5
+ SigmoidGrad 3
+ Slice 2
+ Sum 9
+ TanhGrad 2
+cluster 18 size 22
+ Add 2
+ BiasAdd 1
+ ConcatV2 1
+ Const 3
+ GreaterEqual 1
+ MatMul 1
+ Mul 3
+ Select 3
+ Shape 1
+ Sigmoid 3
+ Split 1
+ Tanh 2
+cluster 19 size 60
+ Add 2
+ AddN 4
+ BiasAddGrad 1
+ BroadcastGradientArgs 1
+ Cast 2
+ ConcatOffset 1
+ ConcatV2 1
+ Const 8
+ MatMul 2
+ Mul 6
+ NoOp 2
+ Reshape 9
+ Select 5
+ SigmoidGrad 3
+ Slice 2
+ Sum 9
+ TanhGrad 2
+cluster 21 size 29
+ Add 2
+ BiasAdd 1
+ Cast 1
+ ConcatV2 1
+ Const 7
+ GreaterEqual 2
+ MatMul 1
+ Mul 5
+ Select 2
+ Sigmoid 3
+ Snapshot 1
+ Split 1
+ Tanh 2
+cluster 22 size 28
+ Add 3
+ BiasAdd 1
+ Cast 1
+ ConcatV2 1
+ Const 5
+ GreaterEqual 1
+ MatMul 1
+ Mul 5
+ Select 3
+ Sigmoid 3
+ Snapshot 1
+ Split 1
+ Tanh 2
+cluster 23 size 423
+ Add 12
+ AddN 28
+ BiasAddGrad 6
+ BroadcastGradientArgs 12
+ Cast 12
+ ConcatOffset 6
+ ConcatV2 6
+ Const 49
+ MatMul 12
+ Mul 48
+ NoOp 13
+ Reshape 76
+ Select 25
+ SigmoidGrad 18
+ Slice 12
+ Sum 76
+ TanhGrad 12
+cluster 24 size 4
+ Const 1
+ Shape 2
+ Transpose 1
+cluster 25 size 27
+ Add 3
+ BiasAdd 1
+ Cast 1
+ ConcatV2 1
+ Const 5
+ GreaterEqual 1
+ MatMul 1
+ Mul 5
+ Select 2
+ Sigmoid 3
+ Snapshot 1
+ Split 1
+ Tanh 2
+cluster 26 size 27
+ Add 3
+ BiasAdd 1
+ Cast 1
+ ConcatV2 1
+ Const 5
+ GreaterEqual 1
+ MatMul 1
+ Mul 5
+ Select 2
+ Sigmoid 3
+ Snapshot 1
+ Split 1
+ Tanh 2
+cluster 27 size 27
+ Add 3
+ BiasAdd 1
+ Cast 1
+ ConcatV2 1
+ Const 5
+ GreaterEqual 1
+ MatMul 1
+ Mul 5
+ Select 2
+ Sigmoid 3
+ Snapshot 1
+ Split 1
+ Tanh 2
+cluster 28 size 27
+ Add 3
+ BiasAdd 1
+ Cast 1
+ ConcatV2 1
+ Const 5
+ GreaterEqual 1
+ MatMul 1
+ Mul 5
+ Select 2
+ Sigmoid 3
+ Snapshot 1
+ Split 1
+ Tanh 2
+cluster 29 size 9
+ Add 1
+ Mul 2
+ RealDiv 2
+ Sqrt 2
+ Sub 2
+cluster 30 size 9
+ Add 1
+ Mul 2
+ RealDiv 2
+ Sqrt 2
+ Sub 2
+cluster 31 size 4
+ Mul 3
+ UnsortedSegmentSum 1
+cluster 32 size 4
+ Mul 3
+ UnsortedSegmentSum 1