}

bool HasTPUDevice(const mlir::TF::RuntimeDevices& devices) {
  for (const auto& device : devices.device_names()) {
    if (device.has_type && device.type == "TPU") return true;
  }
  return false;
}

mlir::LogicalResult GetHostDeviceOutsideCompilationInGenericPipeline(
    mlir::TF::RuntimeDevices devices, std::string* host_device) {
  for (const auto& device : devices.device_names()) {

}

LogicalResult GetTpuDeviceAssignment(
    ClusterOp cluster, ReplicateOp replicate, mlir::ModuleOp module,
    absl::StatusOr<TPUDeviceAssignment>& status_or_tpu_device_assignment) {
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices))) return failure();

  uint32_t num_replicas = replicate.getN();

  auto num_cores_per_replica_attr = cluster->getAttrOfType<mlir::IntegerAttr>(