Scrambles the input key so that users don't need to worry about which part of the key needs to be strong.
| 230 | // Scrambles the input key so that users don't need to worry about which part |
| 231 | // of the key needs to be strong. |
| 232 | std::pair<Philox4x32State, Philox4x32Key> ScramblePhiloxKey(Philox4x32Key key) { |
| 233 | XlaBuilder* builder = key[0].builder(); |
| 234 | XlaOp key0 = ConvertElementType(key[0], U64); |
| 235 | XlaOp key1 = ConvertElementType(key[1], U64); |
| 236 | |
| 237 | Philox4x32State state = { |
| 238 | ConvertElementType(key0, U32), |
| 239 | ConvertElementType(key0 >> ScalarLike(key0, 32), U32), |
| 240 | ConvertElementType(key1, U32), |
| 241 | ConvertElementType(key1 >> ScalarLike(key1, 32), U32), |
| 242 | }; |
| 243 | key = {ConstantR0<uint32>(builder, 0x3ec8f720), |
| 244 | ConstantR0<uint32>(builder, 0x02461e29)}; |
| 245 | state = Philox4x32(state, key); |
| 246 | XlaOp zero = ConstantR0<uint32>(builder, 0); |
| 247 | return {Philox4x32State{zero, zero, state[2], state[3]}, |
| 248 | Philox4x32Key{state[0], state[1]}}; |
| 249 | } |
| 250 | |
| 251 | // Adds an U128 tensor with an U64 tensor. The U128 tensor is represented as two |
| 252 | // U64s with the low 64bits in the front. This routine supports explicit |
no test coverage detected