From 18616396a1356c3da71537cbca3b26009168790c Mon Sep 17 00:00:00 2001 From: Tong Shen Date: Thu, 8 Aug 2019 13:49:15 -0700 Subject: [PATCH] Rollback an accidental change that sneaked into commit 8cc20b. Outcome of that change: for tf.cond(pred, lambda: outside_compilation(A), lambda: outside_compilation(B)), the "Switch" node for input of A will have outside compilation attribute, and we will then mark the whole functionalize If node with outside compialtion attribute. This effectively turns original code into outside_compilation(tf.cond(pred, A, B)), which causes all inputs of A/B to be transferred from TPU to host unconditionally. We should still determine if the whole tf.cond() is in outside compilation by checking whether the predicate is in outside compilation. PiperOrigin-RevId: 262424903 --- tensorflow/compiler/tf2xla/functionalize_cond.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/tf2xla/functionalize_cond.cc b/tensorflow/compiler/tf2xla/functionalize_cond.cc index c709ce30e39..31a0ec794fb 100644 --- a/tensorflow/compiler/tf2xla/functionalize_cond.cc +++ b/tensorflow/compiler/tf2xla/functionalize_cond.cc @@ -787,7 +787,7 @@ Status Conditional::BuildIfNode(Graph* graph, builder.Attr("Tcond", DT_BOOL); string outside_compilation; - if (GetNodeAttr((*switches_.begin())->def(), kXlaOutsideCompilationAttrName, + if (GetNodeAttr(predicate_.node->def(), kXlaOutsideCompilationAttrName, &outside_compilation) .ok()) { builder.Attr(kXlaOutsideCompilationAttrName, outside_compilation);