Remove special TopK handling in SPMD partitioner

PiperOrigin-RevId: 604819001
tensorflow · Feb 7, 2024 · c81766f · c81766f
1 parent 8a81704
commit c81766f
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
@@ -3147,7 +3147,6 @@ bool HloInstruction::has_to_apply() const {
     case HloOpcode::kReduceWindow:
     case HloOpcode::kScatter:
     case HloOpcode::kSort:
-    case HloOpcode::kTopK:
       return true;
     case HloOpcode::kCustomCall:
       // CustomCall can have a to_apply computation, but it is not required to