| 494 | } |
| 495 | |
| 496 | RngOutput UniformIntDistribution(XlaOp key, XlaOp initial_state, |
| 497 | BitGeneratorTy bit_generator, XlaOp minval, |
| 498 | XlaOp maxval, const Shape& shape) { |
| 499 | RngOutput bits_state = bit_generator(key, initial_state, shape); |
| 500 | XlaOp bits = bits_state.value; |
| 501 | XlaOp new_state = bits_state.state; |
| 502 | PrimitiveType type = shape.element_type(); |
| 503 | PrimitiveType unsigned_type; |
| 504 | if (type == U32 || type == S32) { |
| 505 | unsigned_type = U32; |
| 506 | } else { |
| 507 | DCHECK(type == U64 || type == S64); |
| 508 | unsigned_type = U64; |
| 509 | } |
| 510 | return { |
| 511 | ConvertRandomBitsToUniformInt(bits, minval, maxval, type, unsigned_type), |
| 512 | new_state}; |
| 513 | } |
| 514 | |
| 515 | RngOutput NormalFloatingPointDistribution(XlaOp key, XlaOp initial_state, |
| 516 | BitGeneratorTy bit_generator, |
no test coverage detected