From 3b1545a555e14a94a1394d93327d1e6ad64e8319 Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Fri, 24 Apr 2020 16:24:02 -0700
Subject: [PATCH] [Grappler] Fix a bug in folding multiply into add arithmetic
 optimization

PiperOrigin-RevId: 308348068
Change-Id: I61ca637344519f9effee313ea3f911c462d9bfe6
---
 tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index b9502ffb45e..7a9462e88b7 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -2178,6 +2178,8 @@ class FoldMultiplyIntoConv : public ArithmeticOptimizerStage {
     // Check that value preserving chain is the only consumer of the Mul output.
     TF_RETURN_IF_TRUE(!IsAnyMul(*source));
     TF_RETURN_IF_TRUE(NumNonControlOutputs(*source, *ctx().node_map) != 1);
+    // And that Mul is not in the preserve set.
+    TF_RETURN_IF_TRUE(IsInPreserveSet(*source));
 
     const NodeDef* mul = source;
     int input_idx = 0;