Improve performance of domain isolator

Previously we used HloInstruction::ReplaceUseWith what involves a sahpe
check whayt can be expensive for large tuples. After this change we swith
to HloInstruction::ReplaceUseWithDifferentShape what is an equivalent method
just without the shape check to save compile time.

PiperOrigin-RevId: 245729563

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;
         }
       }