| 1394 | } |
| 1395 | |
| 1396 | XlaOp XlaBuilder::InfeedWithToken(XlaOp token, const Shape& shape, |
| 1397 | const string& config) { |
| 1398 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1399 | HloInstructionProto instr; |
| 1400 | if (!LayoutUtil::HasLayout(shape)) { |
| 1401 | return InvalidArgument("Given shape to Infeed must have a layout"); |
| 1402 | } |
| 1403 | const Shape infeed_instruction_shape = |
| 1404 | ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}); |
| 1405 | *instr.mutable_shape() = infeed_instruction_shape.ToProto(); |
| 1406 | instr.set_infeed_config(config); |
| 1407 | |
| 1408 | if (shape.IsArray() && sharding() && |
| 1409 | sharding()->type() == OpSharding::OTHER) { |
| 1410 | // TODO(b/110793772): Support tiled array-shaped infeeds. |
| 1411 | return InvalidArgument( |
| 1412 | "Tiled sharding is not yet supported for array-shaped infeeds"); |
| 1413 | } |
| 1414 | |
| 1415 | if (sharding() && sharding()->type() == OpSharding::REPLICATED) { |
| 1416 | return InvalidArgument( |
| 1417 | "Replicated sharding is not yet supported for infeeds"); |
| 1418 | } |
| 1419 | |
| 1420 | return AddInstruction(std::move(instr), HloOpcode::kInfeed, {token}); |
| 1421 | }); |
| 1422 | } |
| 1423 | |
| 1424 | void XlaBuilder::Outfeed(XlaOp operand, const Shape& shape_with_layout, |
| 1425 | const string& outfeed_config) { |
no test coverage detected