| 1766 | } |
| 1767 | |
| 1768 | XlaOp XlaBuilder::While(const XlaComputation& condition, |
| 1769 | const XlaComputation& body, XlaOp init) { |
| 1770 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1771 | HloInstructionProto instr; |
| 1772 | |
| 1773 | // Infer shape. |
| 1774 | TF_ASSIGN_OR_RETURN(const auto& body_program_shape, body.GetProgramShape()); |
| 1775 | TF_ASSIGN_OR_RETURN(const auto& condition_program_shape, |
| 1776 | condition.GetProgramShape()); |
| 1777 | TF_ASSIGN_OR_RETURN(const Shape* init_shape, GetShapePtr(init)); |
| 1778 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferWhileShape( |
| 1779 | condition_program_shape, |
| 1780 | body_program_shape, *init_shape)); |
| 1781 | *instr.mutable_shape() = shape.ToProto(); |
| 1782 | // Body comes before condition computation in the vector. |
| 1783 | AddCalledComputation(body, &instr); |
| 1784 | AddCalledComputation(condition, &instr); |
| 1785 | return AddInstruction(std::move(instr), HloOpcode::kWhile, {init}); |
| 1786 | }); |
| 1787 | } |
| 1788 | |
| 1789 | XlaOp XlaBuilder::Gather(XlaOp input, XlaOp start_indices, |
| 1790 | const GatherDimensionNumbers& dimension_numbers, |