auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      mlir::UnknownLoc::get(&context), result_types);
  cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));

  mlir::TF::RuntimeDevices devices;
  std::string host_device;
  EXPECT_TRUE(mlir::failed(
      GetHostDeviceOutsideComputation(devices, cluster, &host_device)));
}

}

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

// Returns true if the devices list contain any TPU devices
bool HasTPUDevice(const mlir::TF::RuntimeDevices& devices);

// Returns the host device used for outside compilation in generic pipeline.
mlir::LogicalResult GetHostDeviceOutsideCompilationInGenericPipeline(
    mlir::TF::RuntimeDevices devices, std::string* host_device);

// Parses XLA compilation and execution devices from a tf_device.cluster and

// cluster to before the cluster.
mlir::LogicalResult LiftHeadOutsideCompiledOps(
    OpBuilder* builder,
    const mlir::TF::SideEffectAnalysis& side_effect_analysis,
    const mlir::TF::RuntimeDevices& devices, mlir::tf_device::ClusterOp cluster,
    std::string* host_device, bool* cluster_updated) {
  llvm::SmallVector<Operation*, 4> head_outside_compiled_ops =

}

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

using Permutation = SmallVector<int64_t, 4>;

void LayoutAssignmentPass::runOnOperation() {
  func::FuncOp func = getOperation();

  // Get runtime devices information from the closest parent module.
  RuntimeDevices devices;
  if (failed(::tensorflow::GetDevicesFromOp(func->getParentOfType<ModuleOp>(),
                                            &devices)))
    return signalPassFailure();

      new_parallel_execute, builder);
  if (failed(result)) return failure();

  return TF::RemoveSingletonParallelExecuteOp(new_parallel_execute, builder);
}

void TPURewritePass::runOnOperation() {
  TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(getOperation(), &devices)))
    return signalPassFailure();

  // Collect compilation results.