message: Replace computation in kAllReduce correctly
diff: @@ -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:
  