#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/distributed_runtime/coordination/coordination_service_agent.h"

TF_Device* TF_GetDevice(TF_OpKernelContext* ctx) {
  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
  return reinterpret_cast<TF_Device*>(cc_ctx->device());
}

// --------------------------  Resource  ---------------------------------------

                              const char* device_name, void* device_info,
                              TF_Status* status) {
  // Fill in default values for optional functionality.
  if (device.pack == nullptr) {
    device.pack = &DefaultCustomDevicePack;
  }
  auto custom_device = std::make_unique<tensorflow::CustomDeviceAPI>(
      ctx, device, device_info, device_name);