};

Status XlaCpuDeviceFactory::ListPhysicalDevices(std::vector<string>* devices) {
  XlaDeviceFlags* flags = GetXlaDeviceFlags();
  if (!flags->tf_xla_enable_xla_devices && !XlaDevicesCreationRequired()) {
    VLOG(1) << "Not creating XLA devices, tf_xla_enable_xla_devices not set "
               "and XLA device creation not requested";
    return absl::OkStatus();
  }

};

Status XlaGpuDeviceFactory::ListPhysicalDevices(std::vector<string>* devices) {
  XlaDeviceFlags* flags = GetXlaDeviceFlags();
  if (!flags->tf_xla_enable_xla_devices && !XlaDevicesCreationRequired()) {
    VLOG(1) << "Not creating XLA devices, tf_xla_enable_xla_devices not set "
               "and XLA devices creation not required";
    return absl::OkStatus();
  }

  auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
  rollout_config.enabled_for_xla_launch_ = true;
  rollout_config.enabled_for_compile_on_demand_ = true;

  tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
  return true;
}();

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

            "Switch a device into 'on-demand' mode, where instead of "
            "autoclustering ops are compiled one by one just-in-time."),

       Flag("tf_xla_enable_xla_devices",
            &device_flags->tf_xla_enable_xla_devices,
            "Generate XLA_* devices, where placing a computation on such a "
            "device"
            "forces compilation by XLA. Deprecated."),

    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;

class XlaPlatformInfoTest : public ::testing::Test {
 protected:
  void SetUp() override {
    tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
    tensorflow::GetMarkForCompilationPassFlags()
        ->tf_xla_persistent_cache_directory = "";
    tensorflow::GetMarkForCompilationPassFlags()
        ->tf_xla_persistent_cache_device_types = "";
  }

    rollout_config.enabled_for_gpu_ = true;

    // Set flag to enable using XLA devices. PJRT currently is only supported
    // for XLA devices.
    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);

}

class XlaCompilerOptionsTest : public ::testing::Test {
 protected:
  void SetUp() override {
    tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
  }

  DeviceSetup device_setup_;
};

TEST_F(XlaCompilerOptionsTest, PjRtOptionsXlaDevice) {
  device_setup_.AddDevicesAndSetUp({DEVICE_XLA_GPU});

  // feature is battle-tested, we will switch this to be a session option.
  bool tf_xla_compile_on_demand;

  // Enables "XLA" devices if this flag is set.
  bool tf_xla_enable_xla_devices;
};

// Flags common to the _Xla* ops and their kernels.
struct XlaOpsCommonFlags {
  // If true, _XlaCompile always refuses to compile the cluster, which means the

    rollout_config.enabled_for_compile_on_demand_ = true;

    // Set flag to enable using XLA devices. PJRT currently is only supported
    // 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);

#include "tensorflow/core/platform/test.h"

using ::tensorflow::testing::FindNodeByName;

namespace tensorflow {
namespace {

static bool Initialized = [] {
  tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
  return true;
}();

REGISTER_OP("UncompilableNullary").Output("o: float");
REGISTER_OP("UncompilableUnary").Input("a: float").Output("o: float");