absl::StrCat("found ", kDeviceTPUSystem, " devices with conflicting ",
                   attribute, "s '", a, "' and '", b, "'"));
}

// Find TPU_SYSTEM:0 devices in `devices`. If multiple TPU_SYSTEM devices are
// found, the first one lexicographically is returned. If no TPU_SYSTEM device
// is found or if there are multiple TPU_SYSTEM devices with different jobs or
// replicas, a failure will be returned.

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

// loop that iterates the dataset and calls the reduction function.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeReduceDatasetPass();

// Create a pass that convert ops that copy tensors between devices, e.g.
// tf.Identity.
std::unique_ptr<OperationPass<mlir::func::FuncOp>>
CreateTensorDeviceCopyConversionPass();

// Returns a pass that folds tf.BroadcastTo nodes with subsequent nodes if they