| 1065 | } |
| 1066 | |
| 1067 | XlaOp XlaBuilder::Select(XlaOp pred, XlaOp on_true, XlaOp on_false) { |
| 1068 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1069 | TF_ASSIGN_OR_RETURN(const Shape* true_shape, GetShapePtr(on_true)); |
| 1070 | TF_ASSIGN_OR_RETURN(const Shape* false_shape, GetShapePtr(on_false)); |
| 1071 | TF_RET_CHECK(true_shape->IsTuple() == false_shape->IsTuple()); |
| 1072 | HloOpcode opcode = |
| 1073 | true_shape->IsTuple() ? HloOpcode::kTupleSelect : HloOpcode::kSelect; |
| 1074 | return TernaryOp(opcode, pred, on_true, on_false); |
| 1075 | }); |
| 1076 | } |
| 1077 | |
| 1078 | XlaOp XlaBuilder::Tuple(absl::Span<const XlaOp> elements) { |
| 1079 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |