}
}

ParallelDevice::ParallelDevice(const std::vector<std::string>& devices,
                               bool is_async, int in_flight_nodes_limit)
    : underlying_devices_(devices),
      default_cancellation_manager_(absl::make_unique<CancellationManager>()) {
  device_threads_.reserve(devices.size());
  for (int device_index = 0; device_index < devices.size(); ++device_index) {
    device_threads_.emplace_back(new DeviceThread(

    TF_Status status;
    // Let this custom device choose the device to pin this op on if it
    // implements the pinning function.
    if (device_.shall_pin_to_this_device != nullptr) {
      return device_.shall_pin_to_this_device(tensorflow::wrap(op), &status);
    }
    return errors::Unimplemented("No custom device pinning implementation.");
  }

 private:
  TFE_Context* context_;

 private:
  ParallelTensor(const ParallelDevice& device,
                 std::vector<TensorHandlePtr> tensors,
                 absl::Span<const int64_t> shape, const TF_DataType dtype)
      : device_(device),
        tensors_(std::move(tensors)),
        shape_(std::vector<int64_t>(shape.begin(), shape.end())),
        dtype_(dtype) {}
  ParallelTensor(const ParallelDevice& device,

    }
    auto tag = tensorflow::strings::StrCat("Device #", i, " (", name, ")");
    // Copy to device
    TFE_TensorHandle* hdevice =
        TFE_TensorHandleCopyToDevice(hcpu, ctx, name.c_str(), status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      ADD_FAILURE() << tag << " -- " << TF_Message(status.get());
      continue;
    }
    // Copy from device to the same device.
    TFE_TensorHandle* hdevice2 =

  ASSERT_TRUE(executed);
  TFE_DeleteTensorHandle(retval);

  // Custom device: inputs in different custom devices fails.
  matmul.reset(MatMulOp(context.get(), hcustom0.get(), hcustom1.get()));
  num_retvals = 1;
  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
  ASSERT_NE(TF_OK, TF_GetCode(status.get()));
  ASSERT_TRUE(absl::StrContains(TF_Message(status.get()), custom0));

//
// device_info is an opaque caller-defined type stored with the custom device
// which is passed to the functions referenced in the TFE_CustomDevice struct
// `device` (execute, delete_device, etc.). It can for example contain the
// names of wrapped devices.
//
// There are currently no graph semantics implemented for registered custom
// devices, so executing tf.functions which contain operations placed on the

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

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

  virtual bool UsesTFRT() = 0;

  // List attributes of available devices
  virtual void ListDevices(std::vector<DeviceAttributes>* devices) = 0;

  // Add `devices` into context's device manager. Context's device manager
  // will take ownership and maintain devices' lifetime.
  virtual absl::Status AddDevices(
      std::vector<std::unique_ptr<Device>> devices) = 0;

  // Block until all pending nodes are finished.

  // in an initialized context.
  for (auto d = devices.begin(); d != devices.end();) {
    if (absl::StrContains(d->get()->name(), "CPU:0")) {
      d = devices.erase(d);
    } else {
      ++d;
    }
  }

  status->status = tensorflow::unwrap(ctx)->AddDevices(std::move(devices));
}

void TFE_InsertConfigKeyValue(TFE_Context* ctx, const char* key,
                              const char* value, TF_Status* status) {

      TFE_ContextListDevices(context.get(), status.get()), TF_DeleteDeviceList);
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
  bool has_tpu = false;
  for (int device_index = 0; device_index < TF_DeviceListCount(devices.get());
       ++device_index) {
    std::string device_type =
        TF_DeviceListType(devices.get(), device_index, status.get());