diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-09-16 13:38:24 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-16 13:42:25 -0700 |
| commit | a6ee64cd216b3ac440262e1f4ec7872fe7026df6 (patch) | |
| tree | 52f0bcfc458a17f04093e8420f76a0bc7303a496 /tensorflow/compiler/xla/service/hlo_computation.cc | |
| parent | 92c31bb620b0f8dd6590380dc6a5674f591ce1cb (diff) | |
Conditionally allow changing a non-fusion computation root_instruction shape.
PiperOrigin-RevId: 213191899
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_computation.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_computation.cc | 6 |
1 files changed, 3 insertions, 3 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index 8c6903d766..601a008d9f 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -279,11 +279,11 @@ Status HloComputation::RemoveInstruction(HloInstruction* instruction) { return Status::OK(); } -void HloComputation::set_root_instruction( - HloInstruction* new_root_instruction) { +void HloComputation::set_root_instruction(HloInstruction* new_root_instruction, + bool accept_different_shape) { // The shape of the root (ignoring layout) is an invariant of the computation // for non-fusion cases. - if (!IsFusionComputation()) { + if (!IsFusionComputation() && !accept_different_shape) { CHECK(ShapeUtil::Compatible(new_root_instruction->shape(), root_instruction_->shape())) << new_root_instruction->shape() << " is incompatible with " |