| 434 | } // namespace |
| 435 | |
| 436 | RngOutput ThreeFryBitGenerator(XlaOp key, XlaOp initial_state, |
| 437 | const Shape& shape) { |
| 438 | PrimitiveType type = shape.element_type(); |
| 439 | switch (type) { |
| 440 | case F32: |
| 441 | case U32: |
| 442 | case S32: |
| 443 | return ThreeFryRngBit32(key, initial_state, shape); |
| 444 | case F64: |
| 445 | case U64: |
| 446 | case S64: |
| 447 | return ThreeFryRngBit64(key, initial_state, shape); |
| 448 | default: |
| 449 | return {key.builder()->ReportError(Unimplemented( |
| 450 | "Types other than F32, F64, U32, S32, U64 and S64 " |
| 451 | "are not implemented by ThreeFryBitGenerator; got %s", |
| 452 | primitive_util::LowercasePrimitiveTypeName(type))), |
| 453 | initial_state}; |
| 454 | } |
| 455 | } |
| 456 | |
| 457 | RngOutput PhiloxBitGenerator(XlaOp key, XlaOp initial_state, |
| 458 | const Shape& shape) { |
nothing calls this directly
no test coverage detected