(algorithm: RandomAlgorithm)
| 14 | |
| 15 | |
| 16 | def _rng_algorithm(algorithm: RandomAlgorithm): |
| 17 | assert algorithm == RandomAlgorithm.RNG_THREE_FRY |
| 18 | if algorithm == RandomAlgorithm.RNG_THREE_FRY: |
| 19 | return hlo.RngAlgorithmAttr.get("THREE_FRY") |
| 20 | elif algorithm == RandomAlgorithm.RNG_PHILOX: |
| 21 | return hlo.RngAlgorithmAttr.get("PHILOX") |
| 22 | elif algorithm == RandomAlgorithm.RNG_DEFAULT: |
| 23 | return hlo.RngAlgorithmAttr.get("DEFAULT") |
| 24 | else: |
| 25 | assert False |
| 26 | |
| 27 | |
| 28 | def rng_uint_generator( |
no test coverage detected