Do not hoist nodes that modify frame info.

PiperOrigin-RevId: 195322927

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index 29f4907..adfae2e 100644 (file)
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -1541,7 +1541,7 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
                          const ChainLinkSet& ops) const {
     if (ops.empty()) return true;
     const NodeDef* op0 = ops.begin()->node;
-    if (!IsUnaryElementWise(*op0)) return false;
+    if (ModifiesFrameInfo(*op0) || !IsUnaryElementWise(*op0)) return false;
     for (const auto& link : ops) {
       const NodeDef* op = link.node;
       if (op->device() != root_node.device() || op->op() != op0->op() ||