BasicTestsForTwoDevices(context.get(),
                          "/job:worker/replica:0/task:1/device:CPU:0",
                          "/job:worker/replica:0/task:2/device:CPU:0");

  worker_server1.release();
  worker_server2.release();
}

TEST(PARALLEL_DEVICE, TestAsyncCopyOff) {
  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
      TFE_NewContextOptions(), TFE_DeleteContextOptions);

  return wrap(t);
}

void TF_DeleteExecutionContext(TF_ExecutionContext* c) { unwrap(c)->Release(); }

TF_AbstractOp* TF_NewAbstractOp(TF_ExecutionContext* c) {
  return wrap((unwrap(c)->CreateOperation()));
}

void TF_DeleteAbstractOp(TF_AbstractOp* op) { unwrap(op)->Release(); }

void TF_DeleteAbstractTensor(TF_AbstractTensor* t) { unwrap(t)->Unref(); }

  TF_RETURN_IF_ERROR(ops::Exp(ctx, inputs[0], &exp_output, "Exp"));
  std::unique_ptr<GradientFunction> gradient_function(
      new PassThroughGradientFunction);
  tape.RecordOperation(inputs, {exp_output}, gradient_function.release());
  TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx,
                                          /*targets*/ {exp_output},
                                          /*sources=*/inputs,

  ::tensorflow::Status s =
      ::tensorflow::Env::Default()->NewWritableFile(filename, &f);
  ::tensorflow::Set_TF_Status_from_Status(status, s);

  if (s.ok()) {
    *handle = reinterpret_cast<TF_WritableFileHandle*>(f.release());
  }
}

void TF_CloseWritableFile(TF_WritableFileHandle* handle, TF_Status* status) {
  auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
  TF_SetStatus(status, TF_OK, "");

  int num_dims_out = TF_NumDims(model_out_tensor);
  TF_DeleteTensor(model_out_tensor);

  // If the output is a scalar, then return the scalar output
  if (num_dims_out == 0) {
    outputs[0] = model_out.release();
    return absl::OkStatus();
  }

  // Else, reduce sum the output to get a scalar

  // Will sum all dimensions, so get a Tensor containing [0,...,num_dims_out-1].
  AbstractTensorHandlePtr sum_dims;
  {