experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[TF:XLA] Fix failure in functionalize_control_flow rewrite for Enter nodes that are unused. Make sure we ignore such nodes without producing an error.

Browse Source 
PiperOrigin-RevId: 161136545
This commit is contained in:
Peter Hawkins 2017-07-06 15:07:04 -07:00 committed by TensorFlower Gardener
parent dabcb60bc9
commit 97cbcac453
2 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/compiler/tf2xla/functionalize_control_flow.cc
View File

@ -261,7 +261,10 @@ Status FunctionalizeLoop(Graph* graph, Frame* frame,
std::vector<const Edge*> edges(arg.enter->out_edges().begin(),
arg.enter->out_edges().end());
for (int i = 0; i < edges.size(); ++i) {
TF_RET_CHECK(!edges[i]->IsControlEdge());
if (edges[i]->IsControlEdge() && edges[i]->dst()->IsSink()) {
continue;
}
TF_RET_CHECK(!edges[i]->IsControlEdge()) << edges[i]->src()->name();
Arg new_arg;
new_arg.is_loop_invariant = false;
if (i == 0) {

3
tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc
View File

@ -65,6 +65,9 @@ TEST(FunctionalizeControlFlow, OneLoopVar) {
auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32);
auto enter =
ops::internal::Enter(scope.WithOpName("while/Enter"), source, "aloop");
// Add an unused Enter node. These should be ignored.
auto enter2 =
ops::internal::Enter(scope.WithOpName("while/Enter2"), source, "aloop");
auto merge = ops::Merge(scope.WithOpName("while/Merge"),
std::initializer_list<Input>{enter, dummy});
auto ten = ops::Const<int32>(