#include "xla/tsl/c/tsl_status_internal.h"
#include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
#include "tensorflow/core/common_runtime/composite_device.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/eager/eager_operation.h"
#include "tensorflow/core/distributed_runtime/coordination/coordination_service_error_util.h"
#include "tensorflow/core/framework/function.h"

                              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);

                      send_device.size());
  TFE_OpSetAttrString(op, "recv_device", recv_device.c_str(),
                      recv_device.size());
  TFE_OpSetAttrInt(op, "send_device_incarnation", send_device_incarnation);

  return op;
}

bool GetDeviceName(TFE_Context* ctx, string* device_name,
                   const char* device_type) {

    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  } else if (!async) {
    // Set the local device to CPU to easily validate mirroring
    string cpu_device_name;
    ASSERT_TRUE(GetDeviceName(ctx, &cpu_device_name, "CPU"));
    TFE_OpSetDevice(matmul, cpu_device_name.c_str(), status);
    EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
    auto remote_arg =

  void Release() override { delete this; }
  Status Reset(const char* op, const char* raw_device_name) override {
    if (op_) {
      return errors::FailedPrecondition("Reset called on already built op.");
    }
    if (raw_device_name) {
      device_name_ = raw_device_name;
    }
    op_type_ = op;
    return absl::OkStatus();
  }
  Status SetOpName(const char* const op_name) override {

  TF_Status* status = TF_NewStatus();
  string lib_path =
      tensorflow::GetDataDependencyFilepath(tensorflow::io::JoinPath(
          "tensorflow", "core", "common_runtime", "next_pluggable_device", "c",
          "test_next_pluggable_device_plugin.so"));
  TF_Library* lib = TF_LoadPluggableDeviceLibrary(lib_path.c_str(), status);
  TF_Code code = TF_GetCode(status);
  string status_msg(TF_Message(status));
  TF_DeleteStatus(status);

        output->Unref();
      }
      TF_RETURN_IF_ERROR(ctx->RegisterFunction(func));
    }

    AbstractOperationPtr fn_op(ctx->CreateOperation());
    TF_RETURN_IF_ERROR(fn_op->Reset(fn_name, /*raw_device_name=*/nullptr));
    for (auto input : inputs) {
      TF_RETURN_IF_ERROR(fn_op->AddInput(input));
    }
    int retvals = outputs.size() - null_indices.size();
    std::vector<AbstractTensorHandle*> fn_outputs(retvals);

    TF_OperationDescription* desc =
        TF_NewOperation(host_graph_, func_name_, func_node_name_);
    for (auto input : inputs) {
      TF_AddInput(desc, input);
    }
    // Set device to CPU because some ops inside the function might not be
    // available on GPU.
    TF_SetDevice(desc, "/cpu:0");
    *op = TF_FinishOperation(desc, s_);
    ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);

                            const char* const* underlying_devices,
                            int num_underlying_devices,
                            TFE_CustomDevice* device, void** device_info) {
  device->copy_tensor_to_device = &CopyToParallelDevice;
  device->copy_tensor_from_device = &CopyTensorFromParallelDevice;
  device->delete_device = &DeleteParallelDevice;
  device->execute = &ParallelDeviceExecute;

  TFE_CustomDevice custom_device;
  custom_device.copy_tensor_to_device = &CopyToLoggingDevice;
  custom_device.copy_tensor_from_device = &CopyTensorFromLoggingDevice;
  custom_device.delete_device = &DeleteLoggingDevice;
  custom_device.execute = &LoggingDeviceExecute;
  LoggingDevice* device = new LoggingDevice;
  device->arrived_flag = arrived_flag;
  device->executed_flag = executed_flag;