cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));

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

TEST(TPURewriteDeviceUtilTest, TestGetHostFailDeviceMissingAttributes) {
  mlir::MLIRContext context;
  context.loadDialect<mlir::tf_device::TensorFlowDeviceDialect>();

  auto result = module.walk([&](tf_device::ClusterOp tpu_cluster) {
    std::optional<std::string> host_device;
    if (pass_host_device) {
      std::string host_device_value;
      if (failed(tensorflow::GetHostDeviceOutsideComputation(
              devices, tpu_cluster, &host_device_value)))
        return WalkResult::interrupt();
      host_device = host_device_value;
    }
    cluster_to_host[tpu_cluster.getOperation()] = host_device;

      FindOutsideCompiledOpsAtHead(side_effect_analysis, cluster);
  if (head_outside_compiled_ops.empty()) return mlir::success();
  if (failed(tensorflow::GetHostDeviceOutsideComputation(devices, cluster,
                                                         host_device)))
    return mlir::failure();

  CreateHeadComputation(builder, cluster, head_outside_compiled_ops,

// returns the host device for the head and tail computations. For TPU device,
// if the computation is replicated, GetDeviceAliasForHostOfLogicalCore(0) is
// returned instead.
mlir::LogicalResult GetHostDeviceOutsideComputation(
    mlir::TF::RuntimeDevices devices, mlir::tf_device::ClusterOp cluster,
    std::string* host_device);

// Checks if a device string is a TPU device.
bool IsTPUDevice(llvm::StringRef device);

    }
  }
  if (host_device->empty()) {
    LOG(ERROR) << "Did not find any CPU:0 host devices";
    return mlir::failure();
  }
  return mlir::success();
}

mlir::LogicalResult GetHostDeviceOutsideComputation(
    mlir::TF::RuntimeDevices devices, mlir::tf_device::ClusterOp cluster,
    std::string* host_device) {
  if (HasTPUDevice(devices) ||
      cluster->getParentOfType<mlir::tf_device::ReplicateOp>()) {