diff options
| author |  Sanjoy Das <sanjoy@google.com> | 2018-01-05 10:55:27 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-01-05 10:59:03 -0800 |
| commit | 604272261e37d0bad102d4e4cbf9e90c66489222 (patch) | |
| tree | 90bf3fa17cf2a5451cc84916411d09ad4966f7c0 /tensorflow/compiler/xla/service/while_loop_simplifier.cc | |
| parent | 003ae0eab1d9a95a85e07887cdecc3e8c836b6e3 (diff) | |
[TF:XLA] Correctly simplify while loops with a non-tuple root body
Explicitly bail out if the root instruction of a loop body isn't a tuple()
instruction.
PiperOrigin-RevId: 180948724
Diffstat (limited to 'tensorflow/compiler/xla/service/while_loop_simplifier.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/while_loop_simplifier.cc | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.cc b/tensorflow/compiler/xla/service/while_loop_simplifier.cc index fb0e6f7ce0..1073cc7700 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.cc @@ -306,6 +306,11 @@ static StatusOr<bool> TryRemoveDeadWhileParams(HloInstruction* while_op) { return false; } + if (while_body_root->opcode() != HloOpcode::kTuple) { + VLOG(2) << "While body's root is not a tuple(...) instruction."; + return false; + } + auto print_no_metadata = HloPrintOptions().set_print_metadata(false); // Bail if param0 of while_cond or while_body has users which aren't of type |