message: Add Sort op to the list of ops considered during computation removal
diff: @@ -131,12 +131,14 @@ void HloModule::ReplaceComputations(
  for (std::unique_ptr<HloComputation>& computation : computations_) {
  for (auto* instruction : computation->instructions()) {
  +
  switch (instruction->opcode()) {
  case HloOpcode::kCall:
  case HloOpcode::kMap:
  case HloOpcode::kReduce:
  case HloOpcode::kReduceWindow:
  - case HloOpcode::kScatter: {
  + case HloOpcode::kScatter:
  + case HloOpcode::kSort: {
  HloComputation* new_arg = tensorflow::gtl::FindWithDefault(
  replacements, instruction->to_apply(), nullptr);
  if (new_arg != nullptr) {
  