diff --git a/tensorflow/compiler/xla/service/hlo_domain_isolator.cc b/tensorflow/compiler/xla/service/hlo_domain_isolator.cc index 3746fbbda02..5b388bc0bd8 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_isolator.cc +++ b/tensorflow/compiler/xla/service/hlo_domain_isolator.cc @@ -47,7 +47,11 @@ StatusOr<bool> RunInternal(HloModule* module, HloInstruction* domain = (*creator)(instruction, root, operand); if (domain != nullptr) { VLOG(4) << "New domain: " << domain->ToString(); - TF_RETURN_IF_ERROR(operand->ReplaceUseWith(instruction, domain)); + // Call ReplaceUseWithDifferentShape even though the shapes are + // expected to match to avoid an expensive shape check between the + // original and the new instruction. + TF_RETURN_IF_ERROR( + operand->ReplaceUseWithDifferentShape(instruction, domain)); ++added_domains; } }