message: [XLA:tooling] Fake replay data generated on the host if sufficiently small.
  Prevents us from compiling/running a bunch of device computations that are
  easily elided by using the host to generate the fake data, which de-clutters
  logs and makes dropping into a debugger easier.
diff: @@ -28,8 +28,9 @@ limitations under the License.
  #include "tensorflow/core/platform/types.h"
  namespace xla {
  +namespace {
  -std::unique_ptr<GlobalData> MakeFakeDataOrDie(const Shape& shape,
  +std::unique_ptr<GlobalData> MakeFakeDataViaDeviceOrDie(const Shape& shape,
  Client* client) {
  ComputationBuilder b(
  client,
  @@ -45,6 +46,46 @@ std::unique_ptr<GlobalData> MakeFakeDataOrDie(const Shape& shape,
  .ConsumeValueOrDie();
  }
  +} // namespace
  +
  +std::unique_ptr<GlobalData> MakeFakeDataOrDie(const Shape& shape,
  + Client* client) {
  + if (ShapeUtil::ByteSizeOf(shape) < (1LL << 30)) {
  + std::unique_ptr<Literal> literal = Literal::CreateFromShape(shape);
  + std::minstd_rand0 engine;
  + switch (shape.element_type()) {
  + case F32: {
  + std::uniform_real_distribution<float> generator(0.0f, 1.0f);
  + TF_CHECK_OK(literal->Populate<float>(
  + [&](tensorflow::gtl::ArraySlice<int64> /*indices*/) {
  + return generator(engine);
  + }));
  + break;
  + }
  + case S32: {
  + std::uniform_int_distribution<int32> generator(
  + std::numeric_limits<int32>::lowest(),
  + std::numeric_limits<int32>::max());
  + TF_CHECK_OK(literal->Populate<int32>(
  + [&](tensorflow::gtl::ArraySlice<int64> /*indices*/) {
  + return generator(engine);
  + }));
  + break;
  + }
  + default:
  + LOG(WARNING)
  + << "Unsupported type for host-side fake data generation: "
  + << ShapeUtil::HumanString(shape)
  + << "; falling back to making small amount of fake data via device.";
  + return MakeFakeDataViaDeviceOrDie(shape, client);
  + }
  + return client->TransferToServer(*literal).ValueOrDie();
  + }
  +
  + // If the data is large, generate it on-device.
  + return MakeFakeDataViaDeviceOrDie(shape, client);
  +}
  +
  std::vector<std::unique_ptr<GlobalData>> MakeFakeArgumentsOrDie(
  const Computation& computation, Client* client) {
  auto program_shape =
  