Generates an array of primitive type U32 with the given shape containing random bits generated by the Philox algorithm. Returns the array and the new state of the random number generator.
| 321 | // random bits generated by the Philox algorithm. Returns the array and the new |
| 322 | // state of the random number generator. |
| 323 | RngOutput PhiloxRngBit32(XlaOp op_key, XlaOp initial_state, |
| 324 | const Shape& shape) { |
| 325 | XlaBuilder* builder = op_key.builder(); |
| 326 | const int64 num_elems = ShapeUtil::ElementsIn(shape); |
| 327 | |
| 328 | Philox4x32Key key = Uint64ToUint32s(op_key); |
| 329 | Philox4x32State bits; |
| 330 | XlaOp new_state; |
| 331 | std::tie(bits, new_state) = GeneratePhiloxBits(num_elems, initial_state, key); |
| 332 | // Combining bits[i] in a round-robin fashion, to align with non-XLA |
| 333 | // implementations |
| 334 | int64 bits_len = (num_elems + 3) / 4; |
| 335 | for (auto i = 0; i < 4; ++i) { |
| 336 | bits[i] = Reshape(bits[i], {bits_len, 1}); |
| 337 | } |
| 338 | XlaOp numbers = ConcatInDim(builder, {bits[0], bits[1], bits[2], bits[3]}, |
| 339 | /*dimension=*/1); |
| 340 | numbers = Reshape(numbers, {bits_len * 4}); |
| 341 | numbers = Slice(numbers, /*start_indices=*/{0}, |
| 342 | /*limit_indices=*/{num_elems}, |
| 343 | /*strides=*/{1}); |
| 344 | return {Reshape(numbers, AsInt64Slice(shape.dimensions())), new_state}; |
| 345 | } |
| 346 | |
| 347 | // Generates an array of primitive type U64 with the given shape containing |
| 348 | // random bits generated by the Philox algorithm. Returns the array and the new |
no test coverage detected