commit 647d1b5ff999fedf6d474256d58076daff464337
author Jacques Pienaar <jpienaar@google.com> Thu Oct 24 17:40:23 2019 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Thu Oct 24 17:44:30 2019 -0700
tree 69247426121296d81f3c0f490f59141f07ce5c73
parent 8845f2c52c1c7ba8db63a64493630a50f4f4a303

Fix incorrect use of UpdateEdge for control edges during import

PiperOrigin-RevId: 276601997
Change-Id: Ia6af490a5cea3072a44505935ea6ddedd9c2ce05

tensorflow/compiler/mlir/tensorflow/translate/import_model.cc

diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index cee7137..fe14b34 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -462,10 +462,15 @@
 
   while (!input_node->out_edges().empty()) {
     const Edge* oe = *input_node->out_edges().begin();
-    TF_RETURN_IF_ERROR(graph_->UpdateEdge(
-        placeholder_node,
-        oe->src_output() == Graph::kControlSlot ? Graph::kControlSlot : 0,
-        oe->dst(), oe->dst_input()));
+    // UpdateEdge cannot be used with control edges.
+    if (oe->src_output() == Graph::kControlSlot) {
+      graph_->AddControlEdge(placeholder_node, oe->dst());
+      graph_->RemoveControlEdge(oe);
+      continue;
+    }
+
+    TF_RETURN_IF_ERROR(
+        graph_->UpdateEdge(placeholder_node, 0, oe->dst(), oe->dst_input()));
   }
 
   graph_->RemoveNode(input_node);