Merge pull request #27205 from georgepaw:fix_replace_computations

PiperOrigin-RevId: 241204232
tensorflow · Mar 31, 2019 · 4ac9871 · 4ac9871
2 parents 03feb93 + 29d4ca0
commit 4ac9871
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc
@@ -132,6 +132,7 @@ void HloModule::ReplaceComputations(
   for (std::unique_ptr<HloComputation>& computation : computations_) {
     for (auto* instruction : computation->instructions()) {
       switch (instruction->opcode()) {
+        case HloOpcode::kAllReduce:
         case HloOpcode::kCall:
         case HloOpcode::kMap:
         case HloOpcode::kReduce: