CheckPickDeviceHasError(false, {kCPU0, kCPU1, kGPU0});
}

void SimpleRoundTripTestForDeviceSet(int num_devices) {
  jit::DeviceSet device_set;
  jit::DeviceInfoCache device_info_cache;

  std::vector<string> expected_devices, actual_devices;

  for (int i = 0; i < num_devices; i++) {
    string device_name =
        absl::StrCat("/job:localhost/replica:0/task:0/device:XPU:", i);

  TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get());
  CHECK_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());

  const int num_devices = TF_DeviceListCount(devices);
  for (int i = 0; i < num_devices; ++i) {
    const string dev_type(TF_DeviceListType(devices, i, status.get()));
    CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  TF_DeviceList* device_list = list.get();

  CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);

  const int num_devices = TF_DeviceListCount(device_list);
  LOG(INFO) << "There are " << num_devices << " devices.";
  for (int i = 0; i < num_devices; ++i) {
    const char* device_name = TF_DeviceListName(device_list, i, s);
    CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);

then `num_devices` ops will run with shared_name value `c`.  Failure to do so
will cause the graph execution to fail to complete.

input: the input to the reduction
data: the value of the reduction across all `num_devices` devices.
reduction: the reduction operation to perform.
num_devices: The number of devices participating in this reduction.

  (*module_ref)->setAttr("tf.devices", device_dict);

  mlir::TF::RuntimeDevices devices;
  EXPECT_TRUE(mlir::succeeded(GetDevicesFromOp(*module_ref, &devices)));

  ASSERT_EQ(devices.NumDevices(), 1);
  ASSERT_EQ(devices.device_names().size(), 1);
  ASSERT_EQ(DeviceNameUtils::ParsedNameToString(devices.device_names()[0]),
            "/job:worker/replica:0/task:0/device:CPU:0");
}

    return signalPassFailure();

  // If there is no runtime device information and data format is not explicitly
  // forced, there is nothing to do.
  if (devices.NumDevices() == 0 && force_data_format_.empty()) return;

  func.walk([&](LayoutSensitiveInterface layout_sensitive_interface) {
    // Get desired op data format.
    StringRef target_data_format = force_data_format_;