message: Use stable sort only if it is requested.
  In the HloEvaluator implementation of Sort, choose std::stable_sort only if the
  is_stable boolean is set.
diff: @@ -33,6 +33,7 @@ limitations under the License.
  #include "tensorflow/compiler/xla/map_util.h"
  #include "tensorflow/compiler/xla/primitive_util.h"
  #include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h"
  +#include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
  #include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
  #include "tensorflow/compiler/xla/service/hlo_instruction.h"
  #include "tensorflow/compiler/xla/service/hlo_opcode.h"
  @@ -1493,22 +1494,20 @@ Status HloEvaluator::HandleSort(HloInstruction* sort) {
  std::vector<int64> indices_to_sort(sort_dim_elements);
  std::iota(indices_to_sort.begin(), indices_to_sort.end(), 0);
  Status compare_status = Status::OK();
  - std::stable_sort(
  - indices_to_sort.begin(), indices_to_sort.end(),
  - [sort, &compare_status, &embedded_evaluator, &literals_to_sort](
  - int64 a, int64 b) {
  + auto comparator = [sort, &compare_status, &embedded_evaluator,
  + &literals_to_sort](int64 a, int64 b) {
  std::vector<Literal> literals;
  literals.reserve(2 * sort->operand_count());
  for (int64 i = 0; i < sort->operand_count(); ++i) {
  - auto lhs = ExtractFromIndexPositions(
  - literals_to_sort[i], {a}, /*extract_as_scalar=*/true);
  + auto lhs = ExtractFromIndexPositions(literals_to_sort[i], {a},
  + /*extract_as_scalar=*/true);
  if (!lhs.ok()) {
  compare_status = lhs.status();
  return false;
  }
  literals.push_back(std::move(lhs.ValueOrDie()));
  - auto rhs = ExtractFromIndexPositions(
  - literals_to_sort[i], {b}, /*extract_as_scalar=*/true);
  + auto rhs = ExtractFromIndexPositions(literals_to_sort[i], {b},
  + /*extract_as_scalar=*/true);
  if (!rhs.ok()) {
  compare_status = rhs.status();
  return false;
  @@ -1516,8 +1515,7 @@ Status HloEvaluator::HandleSort(HloInstruction* sort) {
  literals.push_back(std::move(rhs.ValueOrDie()));
  }
  std::vector<const Literal*> literal_ptrs;
  - absl::c_transform(
  - literals, std::back_inserter(literal_ptrs),
  + absl::c_transform(literals, std::back_inserter(literal_ptrs),
  [](const Literal& literal) { return &literal; });
  auto computed_result =
  @@ -1530,7 +1528,13 @@ Status HloEvaluator::HandleSort(HloInstruction* sort) {
  return false;
  }
  return computed_result.ValueOrDie().Get<bool>({});
  - });
  + };
  + if (Cast<HloSortInstruction>(sort)->is_stable()) {
  + std::stable_sort(indices_to_sort.begin(), indices_to_sort.end(),
  + comparator);
  + } else {
  + std::sort(indices_to_sort.begin(), indices_to_sort.end(), comparator);
  + }
  if (!compare_status.ok()) {
  return compare_status;
  }
  