// 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(),

    llvm::SmallVector<std::string, 8>& core_to_host);

// Returns the first TPU device, for use in the non-replicated case. The list of
// TPU devices is retrived from `op`'s module ancestor.
mlir::LogicalResult GetNonReplicatedTPU0(mlir::Operation* op,
                                         std::string* tpu0_device);

// Returns the CPU of the first TPU device, for use in the non-replicated case.

          devices, &host_device))) {
    return mlir::failure();
  } else {
    core_to_host.push_back(host_device);
    return mlir::success();
  }
}

mlir::LogicalResult GetNonReplicatedTPU0(mlir::Operation* op,
                                         std::string* tpu0_device) {
  // Fetch the TPU devices.
  mlir::ModuleOp moduleOp = op->getParentOfType<mlir::ModuleOp>();