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>();