| 1843 | } |
| 1844 | |
| 1845 | XlaOp XlaBuilder::Conditional(XlaOp predicate, XlaOp true_operand, |
| 1846 | const XlaComputation& true_computation, |
| 1847 | XlaOp false_operand, |
| 1848 | const XlaComputation& false_computation) { |
| 1849 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1850 | TF_ASSIGN_OR_RETURN(const xla::Shape* shape, GetShapePtr(predicate)); |
| 1851 | |
| 1852 | if (!ShapeUtil::IsScalar(*shape) || shape->element_type() != PRED) { |
| 1853 | return InvalidArgument( |
| 1854 | "Argument to predicated-Conditional is not a scalar of PRED type " |
| 1855 | "(%s).", |
| 1856 | ShapeUtil::HumanString(*shape)); |
| 1857 | } |
| 1858 | // The index of true_computation must be 0 and that of false computation |
| 1859 | // must be 1. |
| 1860 | return ConditionalImpl(predicate, {&true_computation, &false_computation}, |
| 1861 | {true_operand, false_operand}); |
| 1862 | }); |
| 1863 | } |
| 1864 | |
| 1865 | XlaOp XlaBuilder::Conditional( |
| 1866 | XlaOp branch_index, |
no test coverage detected