// Early return for any node with a device that is not a CPU or GPU.
    DeviceNameUtils::ParsedName parsed;
    if (DeviceNameUtils::ParseFullName(n->requested_device(), &parsed)) {
      if (parsed.type != DEVICE_CPU && parsed.type != DEVICE_GPU) {
        return false;
      }
    }
  }
  return true;
}

// Checks if a graph node is marked to be a guaranteed constant.
bool is_guaranteed_constant(const Node& n) {

  const PjRtBaseDevice::Metadata* pjrt_device_metadata = nullptr;
  std::shared_ptr<se::DeviceMemoryAllocator> custom_allocator;

  const std::string& device_type = device_base->device_type();
  if (device_type == DEVICE_CPU) {
    platform_id = se::host::kHostPlatformId;
  } else if (device_type == DEVICE_GPU) {
    auto device = static_cast<Device*>(device_base);
    platform_id = device->tensorflow_accelerator_device_info()

se::Platform::Id XlaPlatformInfoFromDevice(DeviceBase* device_base) {
  auto device = static_cast<Device*>(device_base);
  se::Platform::Id platform_id = nullptr;
  if (device->device_type() == DEVICE_CPU) {
    platform_id = se::host::kHostPlatformId;
  }

  return platform_id;
}

absl::flat_hash_map<int, int> CreateVariableLookup(
    const std::vector<VariableInfo>& variables) {

                                  std::unique_ptr<Graph>* result) {
  std::vector<std::unique_ptr<Device>> devices;
  devices.push_back(FakeDevice::Make(kDeviceName, DEVICE_GPU));
  devices.push_back(FakeDevice::Make(kHostName, DEVICE_CPU));

  std::unique_ptr<DeviceSet> device_set(new DeviceSet());
  for (auto& device : devices) {
    device_set->AddDevice(device.get());
  }

  auto graph = std::make_unique<Graph>(OpRegistry::Global());

           (policy == XlaOpRegistry::AutoclusteringPolicy::kIfEnabledGlobally &&
            global_jit_level_ != OptimizerOptions::OFF) ||
           (device_type.type_string() == DEVICE_CPU &&
            policy ==
                XlaOpRegistry::AutoclusteringPolicy::kIfExplicitlyRequested &&
            cpu_global_jit_);
  }

    TF_ASSERT_OK(InitOp());
  }

#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  const char* device_name_ = tensorflow::DEVICE_GPU;
#else
  const char* device_name_ = tensorflow::DEVICE_CPU;
#endif
};

// Validates that the tensor has shape and type corresponding to
// dims and dtype.
void validate_tensor(TF_Tensor* tensor, int64_t* dims, int64_t num_dims,

REGISTER_OP("TestOpWithSingleKernel")
    .Input("a: float")
    .Input("b: float")
    .Output("o: float");
REGISTER_KERNEL_BUILDER(
    Name("TestOpWithSingleKernel").Device(tensorflow::DEVICE_CPU), DummyKernel);

TEST(TestKernel, TestGetAllRegisteredKernels) {
  TF_Status* status = TF_NewStatus();
  TF_Buffer* kernel_list_buf = TF_GetAllRegisteredKernels(status);