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);

  (*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");
}