| 1656 | } |
| 1657 | |
| 1658 | XlaOp XlaBuilder::Map(absl::Span<const XlaOp> operands, |
| 1659 | const XlaComputation& computation, |
| 1660 | absl::Span<const int64> dimensions, |
| 1661 | absl::Span<const XlaOp> static_operands) { |
| 1662 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1663 | if (!static_operands.empty()) { |
| 1664 | return Unimplemented("static_operands is not supported in Map"); |
| 1665 | } |
| 1666 | |
| 1667 | HloInstructionProto instr; |
| 1668 | std::vector<const Shape*> operand_shape_ptrs; |
| 1669 | TF_ASSIGN_OR_RETURN(const auto& operand_shapes, GetOperandShapes(operands)); |
| 1670 | absl::c_transform(operand_shapes, std::back_inserter(operand_shape_ptrs), |
| 1671 | [](const Shape& shape) { return &shape; }); |
| 1672 | TF_ASSIGN_OR_RETURN(const ProgramShape& called_program_shape, |
| 1673 | computation.GetProgramShape()); |
| 1674 | TF_ASSIGN_OR_RETURN( |
| 1675 | Shape shape, ShapeInference::InferMapShape( |
| 1676 | operand_shape_ptrs, called_program_shape, dimensions)); |
| 1677 | *instr.mutable_shape() = shape.ToProto(); |
| 1678 | |
| 1679 | Shape output_shape(instr.shape()); |
| 1680 | const int64 output_rank = output_shape.rank(); |
| 1681 | AddCalledComputation(computation, &instr); |
| 1682 | std::vector<XlaOp> new_operands(operands.begin(), operands.end()); |
| 1683 | for (XlaOp& new_operand : new_operands) { |
| 1684 | TF_ASSIGN_OR_RETURN(const Shape* shape, GetShapePtr(new_operand)); |
| 1685 | const int64 rank = shape->rank(); |
| 1686 | if (rank != output_rank) { |
| 1687 | TF_ASSIGN_OR_RETURN(new_operand, |
| 1688 | InDimBroadcast(output_shape, new_operand, {})); |
| 1689 | TF_ASSIGN_OR_RETURN(shape, GetShapePtr(new_operand)); |
| 1690 | } |
| 1691 | if (!ShapeUtil::SameDimensions(output_shape, *shape)) { |
| 1692 | TF_ASSIGN_OR_RETURN(new_operand, |
| 1693 | AddBroadcastSequence(output_shape, new_operand)); |
| 1694 | } |
| 1695 | } |
| 1696 | |
| 1697 | return AddInstruction(std::move(instr), HloOpcode::kMap, new_operands); |
| 1698 | }); |
| 1699 | } |
| 1700 | |
| 1701 | XlaOp XlaBuilder::RngOp(RandomDistribution distribution, |
| 1702 | absl::Span<const XlaOp> parameters, |
no test coverage detected