if (otherDevice->device_) {
      return SetDevice(otherDevice->device_);
    }
    return ChangeResult::NoChange;
  }

  void print(raw_ostream &os) const override {
    if (device_) {
      os << "<" << device_ << ">";
    } else {
      os << "<no device>";
    }
  }

  ChangeResult SetDevice(mlir::StringAttr device) {
    bool changed = (device != device_);

  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;
    rollout_config.AllowForDeviceInXlaLaunch(DeviceType(device_type));

#include "tensorflow/core/framework/tensor_testutil.h"
#include "tsl/lib/core/status_test_util.h"

namespace tensorflow {
namespace {

class XlaHostSendRecvDeviceContextTest : public ::testing::Test {
 public:
  void SetDevice(const string& device_type) {
    auto device_factory = DeviceFactory::GetFactory(device_type);
    SessionOptions options;
    std::vector<std::unique_ptr<Device>> devices;
    Status s = device_factory->CreateDevices(

    auto device =
        DeviceFactory::NewDevice(device_type.type_string(), SessionOptions(),
                                 "/job:localhost/replica:0/task:0");
    device_ = device.get();
    SetDevice(device_type, std::move(device));

    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_fns{
        UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
    device_context_ = core::RefCountPtr<DeviceContext>(

    auto device =
        DeviceFactory::NewDevice(device_type.type_string(), SessionOptions(),
                                 "/job:localhost/replica:0/task:0");
    device_ = device.get();
    SetDevice(device_type, std::move(device));

    // Create PjRtClient for XLA_CPU.
    TF_CHECK_OK(SetPjRtClientInTFGlobalResourceManager(
        device_type,

    TF_DeleteStatus(status);

#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
    std::unique_ptr<Device> device(
        DeviceFactory::NewDevice(device_name_, {}, "/job:a/replica:0/task:0"));
    OpsTestBase::SetDevice(DEVICE_GPU, std::move(device));
#endif
    TF_ASSERT_OK(NodeDefBuilder(op_name, op_name).Finalize(node_def()));
    TF_ASSERT_OK(InitOp());
  }

#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM

};

#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TEST_F(XlaPlatformInfoTest, BuildXlaDeviceCompilerXlaDeviceMetadata) {
  device_setup_.AddDevicesAndSetUp({DEVICE_XLA_GPU});

  Device* device = device_setup_.GetDevice(DEVICE_XLA_GPU);
  const XlaDevice::Metadata* metadata = nullptr;
  TF_CHECK_OK(XlaDevice::GetMetadataFromDevice(device, &metadata));
  XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(device);

  }

  DeviceSetup device_setup_;
};

TEST_F(XlaCompilerOptionsTest, PjRtOptionsXlaDevice) {
  device_setup_.AddDevicesAndSetUp({DEVICE_XLA_GPU});
  Device* device = device_setup_.GetDevice(DEVICE_XLA_GPU);
  DeviceType compilation_device_type = DeviceType(DEVICE_GPU_XLA_JIT);

  se::Platform::Id platform_id = nullptr;
  auto xla_device_metadata = CreateXlaDeviceMetadata(compilation_device_type);

TEST(PjrtCompileUtilTest, CompileToPjRtLoadedExecutable) {
  DeviceSetup device_setup;
  TF_ASSERT_OK_AND_ASSIGN(auto fdef, SampleFuntionAddXY("foo"));
  device_setup.AddDevicesAndSetUp({DEVICE_GPU}, fdef);

  Device* device = device_setup.GetDevice(DEVICE_GPU);
  const XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(device);

  NameAttrList function;
  function.set_name("foo");

  ResourceMgr resource_mgr("");

        return WalkResult::advance();
      } else if (device_str_attr.getValue() != launch.getDevice()) {
        return launch.emitOpError()
               << "inner op has conflicting 'device' attribute, "
                  "got '"
               << device_str_attr.getValue() << "' but expected '"
               << launch.getDevice() << "'";
      }
    } else {
      return launch.emitOpError()