context, parallel_device_name.c_str(), t_released->dtype(), t_released,
      handle_methods, status));
}

// For TFE_CustomDevice::copy_tensor_to_device in the parallel device
// registration.
//
// Since this function is used to satisfy the TFE_CustomDevice C API,
// device_info is passed in using a C-style generic. It must always be a
// ParallelDevice.

      ->tensor;
}

void AllocateLoggingDevice(const char* name, bool* arrived_flag,
                           bool* executed_flag, TFE_CustomDevice** device,
                           void** device_info) {
  TFE_CustomDevice* custom_device = new TFE_CustomDevice;
  custom_device->copy_tensor_to_device = &CopyToLoggingDevice;
  custom_device->copy_tensor_from_device = &CopyTensorFromLoggingDevice;

  return const_cast<void*>(
      static_cast<const void*>(tensor->tensor_data().data()));
}

namespace tensorflow {
namespace {
class CustomDeviceAPI : public tensorflow::CustomDevice {
 public:
  CustomDeviceAPI(TFE_Context* context, TFE_CustomDevice device, void* info,
                  string name)
      : context_(context), device_(device), info_(info), name_(name) {}