| 1736 | } |
| 1737 | |
| 1738 | XlaOp XlaBuilder::RngBitGenerator(RandomAlgorithm algorithm, |
| 1739 | XlaOp initial_state, const Shape& shape) { |
| 1740 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1741 | HloInstructionProto instr; |
| 1742 | TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(shape)); |
| 1743 | TF_ASSIGN_OR_RETURN(Shape state_shape, GetShape(initial_state)); |
| 1744 | Shape output_shape = shape; |
| 1745 | switch (output_shape.element_type()) { |
| 1746 | case PrimitiveType::F32: |
| 1747 | case PrimitiveType::S32: |
| 1748 | case PrimitiveType::U32: |
| 1749 | output_shape.set_element_type(PrimitiveType::U32); |
| 1750 | break; |
| 1751 | case PrimitiveType::F64: |
| 1752 | case PrimitiveType::S64: |
| 1753 | case PrimitiveType::U64: |
| 1754 | output_shape.set_element_type(PrimitiveType::U64); |
| 1755 | break; |
| 1756 | default: |
| 1757 | return InvalidArgument("Unsupported shape for RngBitGenerator: %s", |
| 1758 | PrimitiveType_Name(output_shape.element_type())); |
| 1759 | } |
| 1760 | *instr.mutable_shape() = |
| 1761 | ShapeUtil::MakeTupleShape({state_shape, output_shape}).ToProto(); |
| 1762 | instr.set_rng_algorithm(algorithm); |
| 1763 | return AddInstruction(std::move(instr), HloOpcode::kRngBitGenerator, |
| 1764 | {initial_state}); |
| 1765 | }); |
| 1766 | } |
| 1767 | |
| 1768 | XlaOp XlaBuilder::While(const XlaComputation& condition, |
| 1769 | const XlaComputation& body, XlaOp init) { |
no test coverage detected