| 1603 | } |
| 1604 | |
| 1605 | XlaOp XlaBuilder::Sort(absl::Span<const XlaOp> operands, |
| 1606 | const XlaComputation& comparator, int64 dimension, |
| 1607 | bool is_stable) { |
| 1608 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1609 | HloInstructionProto instr; |
| 1610 | instr.set_is_stable(is_stable); |
| 1611 | std::vector<const Shape*> operand_shape_ptrs; |
| 1612 | TF_ASSIGN_OR_RETURN(std::vector<Shape> operand_shapes, |
| 1613 | GetOperandShapes(operands)); |
| 1614 | absl::c_transform(operand_shapes, std::back_inserter(operand_shape_ptrs), |
| 1615 | [](const Shape& shape) { return &shape; }); |
| 1616 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferVariadicOpShape( |
| 1617 | HloOpcode::kSort, operand_shape_ptrs)); |
| 1618 | *instr.mutable_shape() = shape.ToProto(); |
| 1619 | if (dimension == -1) { |
| 1620 | TF_ASSIGN_OR_RETURN(const Shape* keys_shape, GetShapePtr(operands[0])); |
| 1621 | dimension = keys_shape->rank() - 1; |
| 1622 | } |
| 1623 | instr.add_dimensions(dimension); |
| 1624 | AddCalledComputation(comparator, &instr); |
| 1625 | return AddInstruction(std::move(instr), HloOpcode::kSort, operands); |
| 1626 | }); |
| 1627 | } |
| 1628 | |
| 1629 | XlaOp XlaBuilder::ConvertElementType(XlaOp operand, |
| 1630 | PrimitiveType new_element_type) { |
no test coverage detected