TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  // Create a parallel device with two CPUs
  const char* first_device_name =
      "/job:localhost/replica:0/task:0/device:CUSTOM:0";
  std::array<const char*, 2> first_underlying_devices{
      "/job:localhost/replica:0/task:0/device:CPU:0",

void TFE_SetLogicalCpuDevices(TFE_Context* ctx, int num_cpus,
                              const char* prefix, TF_Status* status) {
  std::vector<std::unique_ptr<tensorflow::Device>> devices;

  if (prefix == nullptr || strlen(prefix) == 0)
    prefix = "/job:localhost/replica:0/task:0";

  tensorflow::SessionOptions sess_options;
  (*sess_options.config.mutable_device_count())["CPU"] = num_cpus;
  status->status =