limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_
#define TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_

// A simple logging device to test custom device registration.
#include <memory>

#include "tensorflow/c/c_api.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_experimental.h"

void AllocateParallelDevice(const char* device_name,
                            const char* const* underlying_devices,
                            int num_underlying_devices,
                            TFE_CustomDevice* device, void** device_info);

}  // namespace parallel_device
}  // namespace tensorflow

  const char* first_device = "/job:worker/replica:0/task:1/device:CPU:0";
  const char* second_device = "/job:worker/replica:0/task:2/device:CPU:0";
  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
  std::array<const char*, 2> underlying_devices{first_device, second_device};
  RegisterParallelDevice(context.get(), device_name, underlying_devices,
                         status.get());

}

template <std::size_t num_devices>
void RegisterParallelDevice(
    TFE_Context* context, const char* device_name,
    const std::array<const char*, num_devices>& underlying_devices,
    TF_Status* status) {
  TFE_CustomDevice device;
  void* device_info;
  tensorflow::parallel_device::AllocateParallelDevice(
      device_name, underlying_devices.data(), underlying_devices.size(),
      &device, &device_info);

  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;

# tensorflow/python:_pywrap_parallel_device.
filegroup(
    name = "lib_headers",
    srcs = ["parallel_device_lib.h"],
)

filegroup(
    name = "lib_sources",
    srcs = ["parallel_device_lib.cc"],
)

filegroup(
    name = "device_headers",
    srcs = ["parallel_device.h"],
)

filegroup(
    name = "device_sources",
    srcs = ["parallel_device.cc"],
)

filegroup(
    name = "headers",

  // devices.
  virtual Status Dim(int dim_index, int64_t* dim) const = 0;

  // Returns the device which created the handle.
  virtual const char* DeviceName(Status* status) const = 0;
  // Returns the device where the tensor was placed.
  virtual const char* BackingDeviceName(Status* status) const = 0;
  // Returns the device type which created the handle.

// Set whether to propagate assigned device information when constructing a new
// Graph from a GraphDef. By default assigned device information is not copied
// and is re-computed by the runtime.
inline void TF_ImportGraphDefOptionsSetPropagateDeviceSpec(
    TF_ImportGraphDefOptions* opts, unsigned char propagate_device_spec) {
  opts->opts.propagate_device_spec = propagate_device_spec;
}

}  // end namespace tensorflow

-   type: input id: OS attributes: label: OS Platform and Distribution
    description: placeholder: e.g., Linux Ubuntu 16.04 validations: required:
    false

-   type: input id: Mobile attributes: label: Mobile device description:
    placeholder: e.g., Linux Ubuntu 16.04 validations: required: false

-   type: input id: Python attributes: label: Python version description:
    placeholder: e.g., 3.9 validations: required: false

  // returned by DeviceName should be "/device:GPU:*" until a particular GPU is
  // chosen for the operation by the device placement logic in the
  // executor. After that, the value returned by DeviceName will be a full
  // device name such as "/job:localhost/replica:0/task:0/device:GPU:1".
  virtual const string& DeviceName() const = 0;

  // Sets the operation device name.
  //