| 455 | } |
| 456 | |
| 457 | RngOutput PhiloxBitGenerator(XlaOp key, XlaOp initial_state, |
| 458 | const Shape& shape) { |
| 459 | PrimitiveType type = shape.element_type(); |
| 460 | switch (type) { |
| 461 | case F32: |
| 462 | case U32: |
| 463 | case S32: |
| 464 | return PhiloxRngBit32(key, initial_state, shape); |
| 465 | case F64: |
| 466 | case U64: |
| 467 | case S64: |
| 468 | return PhiloxRngBit64(key, initial_state, shape); |
| 469 | default: |
| 470 | return {key.builder()->ReportError(Unimplemented( |
| 471 | "Types other than F32, F64, U32, S32, U64 and S64 " |
| 472 | "are not implemented by PhiloxFryBitGenerator; got %s", |
| 473 | primitive_util::LowercasePrimitiveTypeName(type))), |
| 474 | initial_state}; |
| 475 | } |
| 476 | } |
| 477 | |
| 478 | std::pair<XlaOp, XlaOp> ScramblePhiloxKey(XlaOp key) { |
| 479 | Philox4x32Key pkey = Uint64ToUint32s(key); |
no test coverage detected