// Set the copy op's device to the first TPU.
  std::string device_str;
  if (replicated) {
    device_str = tensorflow::GetDeviceAliasForLogicalCore(0);
  } else if (failed(tensorflow::GetNonReplicatedTPU0(
                 tpu_copy_with_dynamic_shape_op, &device_str))) {
    return failure();
  }

  new_device_launch_op = builder.create<tf_device::LaunchOp>(
      tpu_copy_with_dynamic_shape_op->getLoc(),