From e73efbcf242222df682e4e7388f120ca6401b89f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 17 May 2018 16:52:31 -0700
Subject: [PATCH] Enable optimizations submitted during pre-NIPS freeze:   *
 Snapshot to Identity promotion.   * Removal of transposes pairs separated by
 chain.

PiperOrigin-RevId: 197073602
---
 .../core/grappler/optimizers/arithmetic_optimizer.cc     |  7 ++-----
 .../grappler/optimizers/arithmetic_optimizer_test.cc     |  2 +-
 tensorflow/core/grappler/optimizers/constant_folding.cc  | 16 +++++-----------
 .../core/grappler/optimizers/constant_folding_test.cc    |  6 ++----
 4 files changed, 10 insertions(+), 21 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index adef75f..e7f70c6 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -1175,11 +1175,8 @@ class RemoveIdentityTranspose : public ArithmeticOptimizerStage {
   Status TrySimplify(NodeDef* node, string* simplified_node_name) override {
     TF_RETURN_IF_ERROR(EnsureNodeIsSupported(node));
     NodeDef* tail = node;
-    // TODO(rmlarsen): Enable in regular mode after May 15, 2018.
-    if (opt_level_ == RewriterConfig::AGGRESSIVE) {
-      tail = GetTailOfIdempotentChain(*tail, *ctx().node_map,
-                                      *ctx().nodes_to_preserve);
-    }
+    tail = GetTailOfIdempotentChain(*tail, *ctx().node_map,
+                                    *ctx().nodes_to_preserve);
     NodeDef* first_transpose;
     TF_RETURN_IF_ERROR(GetInputNode(tail->input(0), &first_transpose));
 
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
index 27c0dde..99f93e6 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
@@ -1329,7 +1329,7 @@ TEST_F(ArithmeticOptimizerTest, RemoveIdentityTransposesThroughChain) {
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
   GraphDef output;
-  ArithmeticOptimizer optimizer(RewriterConfig::AGGRESSIVE);
+  ArithmeticOptimizer optimizer;
   EnableOnlyRemoveIdentityTranspose(&optimizer);
   OptimizeAndPrune(&optimizer, &item, &output);
 
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
index b2dcbf9..782ccff 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -1516,10 +1516,7 @@ void ConstantFolding::ReplaceOperationWithSnapshot(
     GraphDef* graph) {
   // If the graph contains no ops that mutate their inputs, we can
   // use Identity insted of Snapshot.
-
-  // TODO(rmlarsen): Enable in regular mode after May 15, 2018.
-  if (opt_level_ == RewriterConfig::AGGRESSIVE &&
-      !graph_contains_assign_or_inplace_op_) {
+  if (!graph_contains_assign_or_inplace_op_) {
     ReplaceOperationWithIdentity(input_to_forward, properties, node, graph);
     return;
   }
@@ -2571,13 +2568,10 @@ Status ConstantFolding::Optimize(Cluster* cluster, const GrapplerItem& item,
   }
 
   graph_contains_assign_or_inplace_op_ = false;
-  // TODO(rmlarsen): Enable in regular mode after May 15, 2018.
-  if (opt_level_ == RewriterConfig::AGGRESSIVE) {
-    for (const NodeDef& node : item.graph.node()) {
-      if (ModifiesInputsInPlace(node) || MaybeHasRefInput(node)) {
-        graph_contains_assign_or_inplace_op_ = true;
-        break;
-      }
+  for (const NodeDef& node : item.graph.node()) {
+    if (ModifiesInputsInPlace(node) || MaybeHasRefInput(node)) {
+      graph_contains_assign_or_inplace_op_ = true;
+      break;
     }
   }
 
diff --git a/tensorflow/core/grappler/optimizers/constant_folding_test.cc b/tensorflow/core/grappler/optimizers/constant_folding_test.cc
index 0bf51c4..a41e0ab 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding_test.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding_test.cc
@@ -72,8 +72,7 @@ class ConstantFoldingTest : public GrapplerTest {
       GrapplerItem item;
       TF_CHECK_OK(s.ToGraphDef(&item.graph));
       item.fetch = {"mul1", "mul2", "add1", "add2"};
-      ConstantFolding optimizer(RewriterConfig::AGGRESSIVE,
-                                nullptr /* cpu_device */);
+      ConstantFolding optimizer(nullptr /* cpu_device */);
       GraphDef output;
       Status status = optimizer.Optimize(nullptr, item, &output);
       TF_EXPECT_OK(status);
@@ -296,8 +295,7 @@ TEST_F(ConstantFoldingTest, NeutralElement) {
     TF_CHECK_OK(s.ToGraphDef(&item.graph));
     item.fetch = {"stack", "matmul3", "matmul4"};
 
-    ConstantFolding optimizer(RewriterConfig::AGGRESSIVE,
-                              nullptr /* cpu_device */);
+    ConstantFolding optimizer(nullptr /* cpu_device */);
     GraphDef output;
     Status status = optimizer.Optimize(nullptr, item, &output);
     TF_EXPECT_OK(status);
-- 
2.7.4