// Flag is turned on, some device is allowlisted, but the requested one isn't.
  rollout_config.AllowForDeviceInXlaLaunch(DeviceType(DEVICE_GPU));
  EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

  // Flag is turned on and the requested device is allowlisted.
  rollout_config.AllowForDeviceInXlaLaunch(DeviceType(DEVICE_CPU));
  EXPECT_TRUE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

}();

class DeviceContextTest : public ::testing::Test {
 public:
  void SetDevice(const string& device_type) {
    auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
    rollout_config.AllowForDeviceInXlaLaunch(DeviceType(device_type));
    rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(device_type));

    auto device_factory = DeviceFactory::GetFactory(device_type);
    SessionOptions options;

    GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;

    // Add and setup the GPU device.
    auto device_type = DeviceType(DEVICE_GPU);
    auto jit_device_type = DeviceType(DEVICE_GPU);
    rollout_config.AllowForDeviceInXlaLaunch(device_type);
    rollout_config.AllowForDeviceInXlaCompileOnDemand(device_type);

    auto device =
        DeviceFactory::NewDevice(device_type.type_string(), SessionOptions(),

    // to the `device_type` being allowed in order to use the Device API for
    // single device compilation and execution in the XlaLaunch op.
    void AllowForDeviceInXlaLaunch(const DeviceType& device_type) {
      xla_launch_allowed_devices_.insert(device_type.type_string());
    }

    bool IsEnabledInXlaLaunchForDevice(const DeviceType& device_type) const {

    // for XLA devices.
    GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;

    // Add and setup the XLA_CPU device.
    auto device_type = DeviceType(DEVICE_XLA_CPU);
    rollout_config.AllowForDeviceInXlaLaunch(device_type);
    rollout_config.AllowForDeviceInXlaCompileOnDemand(device_type);

    auto jit_device_type = DeviceType(DEVICE_CPU_XLA_JIT);
    auto device =