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

  void AddGpuDevice(const ParsedName& device,
                    const GpuDeviceMetadata& metadata);

  llvm::ArrayRef<ParsedName> device_names() const { return device_names_; }
  size_t NumDevices() const { return device_names_.size(); }

  // Returns GPU device metadata if it is available, otherwise returns None.
  std::optional<GpuDeviceMetadata> GetGpuDeviceMetadata(
      const ParsedName& device) const;

// CHECK: %{{.*}} = tf_executor.island(%[[CT_0]], %[[CT_1]])


// Tests devices are not remapped if no devices were defined in replicate.
// CHECK-LABEL: func @no_devices
func.func @no_devices() {
  tf_executor.graph {
    %0 = tf_executor.island {
      tf_device.replicate {n = 2 : i32} {
        "tf_device.launch"() ({
          "tf.opA"() : () -> ()
          tf_device.return

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