| 1589 | } |
| 1590 | |
| 1591 | XlaOp XlaBuilder::Rev(XlaOp operand, absl::Span<const int64> dimensions) { |
| 1592 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1593 | HloInstructionProto instr; |
| 1594 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 1595 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferReverseShape( |
| 1596 | *operand_shape, dimensions)); |
| 1597 | *instr.mutable_shape() = shape.ToProto(); |
| 1598 | for (int64 dim : dimensions) { |
| 1599 | instr.add_dimensions(dim); |
| 1600 | } |
| 1601 | return AddInstruction(std::move(instr), HloOpcode::kReverse, {operand}); |
| 1602 | }); |
| 1603 | } |
| 1604 | |
| 1605 | XlaOp XlaBuilder::Sort(absl::Span<const XlaOp> operands, |
| 1606 | const XlaComputation& comparator, int64 dimension, |