[XLA][TPU] Fix an XLA crash

This keeps outputs of entry computation when it's fused into another compuatation. PiperOrigin-RevId: 637991294
tensorflow · Jun 11, 2024 · a484212 · a484212
1 parent a415715
commit a484212
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc
@@ -2365,10 +2365,12 @@ void HloFusionInstruction::MergeFusionInstructionIntoMultiOutput(
             break;
           }
         }
-        if (has_outside_user) {
-          new_roots.insert(
-              FindOrDie(old_to_new, old_fusion_outputs.back().first));
+        if (!has_outside_user && !user->IsRoot()) {
+          continue;
         }
+
+        new_roots.insert(
+            FindOrDie(old_to_new, old_fusion_outputs.back().first));
       }
       continue;
     }