Generates an array of primitive type U64 with the given shape containing random bits generated by the Philox algorithm. Returns the array and the new state of the random number generator.
| 348 | // random bits generated by the Philox algorithm. Returns the array and the new |
| 349 | // state of the random number generator. |
| 350 | RngOutput PhiloxRngBit64(XlaOp op_key, XlaOp initial_state, |
| 351 | const Shape& shape) { |
| 352 | XlaBuilder* builder = op_key.builder(); |
| 353 | const int64 num_elems = ShapeUtil::ElementsIn(shape); |
| 354 | |
| 355 | Philox4x32Key key = Uint64ToUint32s(op_key); |
| 356 | Philox4x32State bits32; |
| 357 | XlaOp new_state; |
| 358 | std::tie(bits32, new_state) = |
| 359 | GeneratePhiloxBits(num_elems * 2, initial_state, key); |
| 360 | |
| 361 | std::array<XlaOp, 2> bits64; |
| 362 | bits64[0] = Uint32sToUint64({bits32[0], bits32[1]}); |
| 363 | bits64[1] = Uint32sToUint64({bits32[2], bits32[3]}); |
| 364 | |
| 365 | // Combining bits64[i] in a round-robin fashion, to align with non-XLA |
| 366 | // implementations |
| 367 | int64 bits64_len = (num_elems + 1) / 2; |
| 368 | for (auto i = 0; i < 2; ++i) { |
| 369 | bits64[i] = Reshape(bits64[i], {bits64_len, 1}); |
| 370 | } |
| 371 | XlaOp numbers = ConcatInDim(builder, {bits64[0], bits64[1]}, |
| 372 | /*dimension=*/1); |
| 373 | numbers = Reshape(numbers, {bits64_len * 2}); |
| 374 | numbers = Slice(numbers, /*start_indices=*/{0}, |
| 375 | /*limit_indices=*/{num_elems}, |
| 376 | /*strides=*/{1}); |
| 377 | return {Reshape(numbers, AsInt64Slice(shape.dimensions())), new_state}; |
| 378 | } |
| 379 | |
| 380 | XlaOp ConvertRandomBitsToUniformFloatingPoint(XlaOp bits, XlaOp minval, |
| 381 | XlaOp maxval) { |
no test coverage detected