Function ThreeFryRngBit64

tensorflow/compiler/xla/client/lib/prng.cc:169–178

Generates random 64bits with the given shape using the Three Fry implementation. Returns the random bits and the new state.

// Generates random 64bits with the given shape using the Three Fry
// implementation. Returns the random bits and the new state.
RngOutput ThreeFryRngBit64(XlaOp key, XlaOp initial_state, const Shape& shape) {
  const int64 size = ShapeUtil::ElementsIn(shape);
  std::pair<ThreeFry2x32State, XlaOp> inputs_state =
      GetThreeFryInputsAndUpdatedState(initial_state, size);
  ThreeFry2x32State inputs = inputs_state.first;
  ThreeFry2x32State outputs = ThreeFry2x32(inputs, Uint64ToUint32s(key));
  XlaOp result = Uint32sToUint64(outputs);
  return {Reshape(result, AsInt64Slice(shape.dimensions())),
          inputs_state.second};
}

// The key of the Philox random number generator.
using Philox4x32Key = std::array<XlaOp, 2>;