From a6ee64cd216b3ac440262e1f4ec7872fe7026df6 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sun, 16 Sep 2018 13:38:24 -0700 Subject: Conditionally allow changing a non-fusion computation root_instruction shape. PiperOrigin-RevId: 213191899 --- tensorflow/compiler/xla/service/hlo_computation.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) (limited to 'tensorflow/compiler/xla/service/hlo_computation.cc') 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 " -- cgit v1.2.3