| 398 | } |
| 399 | |
| 400 | StatusOr<HloInstruction*> MakeSortHlo( |
| 401 | const Shape& sort_shape, absl::Span<HloInstruction* const> operands, |
| 402 | int64 dimension_to_sort, bool is_stable, HloComputation::Builder* builder, |
| 403 | HloModule* module) { |
| 404 | CHECK(!operands.empty()) << "Sort Hlo requires at least one operand."; |
| 405 | HloComputation* compare_computation; |
| 406 | XlaBuilder b("Sort.Compare"); |
| 407 | std::vector<PrimitiveType> operand_types(operands.size()); |
| 408 | for (int64 i = 0; i < operands.size(); ++i) { |
| 409 | operand_types[i] = operands[i]->shape().element_type(); |
| 410 | } |
| 411 | XlaComputation comparator = CreateScalarLtComputation(operand_types, &b); |
| 412 | TF_ASSIGN_OR_RETURN(ProgramShape program_shape, comparator.GetProgramShape()); |
| 413 | HloModuleConfig config(program_shape); |
| 414 | TF_ASSIGN_OR_RETURN(auto new_module, |
| 415 | HloModule::CreateFromProto(comparator.proto(), config)); |
| 416 | HloCloneContext context(module); |
| 417 | compare_computation = |
| 418 | module->DeepCloneComputation(new_module->entry_computation(), &context); |
| 419 | return builder->AddInstruction(HloInstruction::CreateSort( |
| 420 | sort_shape, dimension_to_sort, operands, compare_computation, is_stable)); |
| 421 | } |
| 422 | |
| 423 | StatusOr<HloInstruction*> CollapseFirstNDims(HloInstruction* operand, int64 n) { |
| 424 | CHECK_GT(n, 0); |