Check for the value of the _XlaCompile boolean since xla often creates it with

the default value of false.

PiperOrigin-RevId: 189041421

diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
index 21037ff..821b22a 100644 (file)
--- a/tensorflow/core/grappler/optimizers/constant_folding.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -752,7 +752,8 @@ bool ConstantFolding::IsFoldable(const NodeDef& node) const {
   if (op.find("Quantized") != string::npos || op.find("Sparse") == 0) {
     return false;
   }
-  if (node.attr().count("_XlaCompile") > 0) {
+  if (node.attr().count("_XlaCompile") > 0 &&
+      node.attr().at("_XlaCompile").b()) {
     return false;
   }