| 1510 | } |
| 1511 | |
| 1512 | XlaOp XlaBuilder::AfterAll(absl::Span<const XlaOp> tokens) { |
| 1513 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1514 | if (tokens.empty()) { |
| 1515 | return InvalidArgument("AfterAll requires at least one operand"); |
| 1516 | } |
| 1517 | for (int i = 0; i < tokens.size(); ++i) { |
| 1518 | XlaOp operand = tokens[i]; |
| 1519 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 1520 | if (!operand_shape->IsToken()) { |
| 1521 | return InvalidArgument( |
| 1522 | "All operands to AfterAll must be tokens; operand %d has shape %s", |
| 1523 | i, ShapeUtil::HumanString(*operand_shape)); |
| 1524 | } |
| 1525 | } |
| 1526 | HloInstructionProto instr; |
| 1527 | *instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); |
| 1528 | return AddInstruction(std::move(instr), HloOpcode::kAfterAll, tokens); |
| 1529 | }); |
| 1530 | } |
| 1531 | |
| 1532 | XlaOp XlaBuilder::CustomCall( |
| 1533 | const string& call_target_name, absl::Span<const XlaOp> operands, |
no test coverage detected