From: A. Unique TensorFlower <gardener@tensorflow.org>
Date: Wed, 9 May 2018 22:58:28 +0000 (-0700)
Subject: Add missing update of node map in the Mul(x,x) => Square(x) rewrite. This is what... 
X-Git-Tag: upstream/v1.9.0_rc1~116^2^2~178
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=f1badb6664c290176864d1a1d4ab537b7332b730;p=platform%2Fupstream%2Ftensorflow.git

Add missing update of node map in the Mul(x,x) => Square(x) rewrite. This is what caused a failure in //photos/vision/object_detection/ranking:brain_embedder_test when the concat/split hoisting was enabled.

PiperOrigin-RevId: 196043455
---

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index adfae2e..f46c30c 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -2233,6 +2233,9 @@ string ArithmeticOptimizer::TrySimplifyAndReplaceUses(
         new_square_node->set_input(i - 1, new_square_node->input(i));
       }
       new_square_node->mutable_input()->RemoveLast();
+      for (const string& input : new_square_node->input()) {
+        node_map_->AddOutput(NodeName(input), new_square_node->name());
+      }
       return new_square_node->name();
     }
   }