| 38 | namespace { |
| 39 | |
| 40 | xla::BitGeneratorTy BitGen(Algorithm alg) { |
| 41 | if (alg == RNG_ALG_PHILOX) { |
| 42 | return [=](xla::XlaOp key, xla::XlaOp state, const xla::Shape& shape) { |
| 43 | state = |
| 44 | xla::ConcatInDim(key.builder(), {xla::Reshape(key, {1}), state}, 0); |
| 45 | xla::XlaOp result = |
| 46 | xla::RngBitGenerator(xla::RandomAlgorithm::RNG_PHILOX, state, shape); |
| 47 | xla::XlaOp data = xla::GetTupleElement(result, 1); |
| 48 | xla::XlaOp new_state = |
| 49 | xla::Slice(xla::GetTupleElement(result, 0), {1}, {3}, {1}); |
| 50 | return xla::RngOutput{data, new_state}; |
| 51 | }; |
| 52 | } else { |
| 53 | return [=](xla::XlaOp key, xla::XlaOp state, const xla::Shape& shape) { |
| 54 | state = xla::ConcatScalars(key.builder(), {key, state}); |
| 55 | xla::XlaOp result = xla::RngBitGenerator( |
| 56 | xla::RandomAlgorithm::RNG_THREE_FRY, state, shape); |
| 57 | xla::XlaOp data = xla::GetTupleElement(result, 1); |
| 58 | xla::XlaOp new_state = xla::Reshape( |
| 59 | xla::Slice(xla::GetTupleElement(result, 0), {1}, {2}, {1}), {}); |
| 60 | return xla::RngOutput{data, new_state}; |
| 61 | }; |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | xla::RngOutput StatefulRngUniform(Algorithm alg, xla::XlaOp key, |
| 66 | xla::XlaOp initial_state, |
no test coverage detected