device_cluster.erase();

  return mlir::success();
}

// Check that cluster results are valid. An result is invalid when it does not
// have a valid XLA type.
LogicalResult CheckClusterResults(mlir::tf_device::ClusterOp cluster) {
  for (OpResult result : cluster.getResults()) {
    if (!tensorflow::TypeValidForXLA(result.getType())) {
      return cluster.emitError()