DefineFunction("MyGrad2", &grad_func2);

  // func cannot be null
  TF_GraphCopyFunction(host_graph_, nullptr, func_, s_);
  EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_));
  EXPECT_EQ(string("'func' argument to TF_GraphCopyFunction cannot be null"),
            string(TF_Message(s_)));

  // Cannot change gradient
  TF_GraphCopyFunction(host_graph_, func_, grad_func1, s_);

  return func->record->fdef().signature().name().c_str();
}

void TF_GraphCopyFunction(TF_Graph* g, const TF_Function* func,
                          const TF_Function* grad, TF_Status* status) {
  if (func == nullptr) {
    status->status = InvalidArgument(
        "'func' argument to TF_GraphCopyFunction cannot be null");
    return;
  }

  tensorflow::mutex_lock l(g->mu);

  TF_Status* s_;
  TF_Graph* func_graph_;
  TF_Function* func_;
};

TEST_F(CApiExperimentalFunctionTest, GraphRemoveFunction) {
  TF_Function* funcs[1];
  DefineFunction(func_name_, &func_);

  TF_GraphCopyFunction(func_graph_, func_, nullptr, s_);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);

  EXPECT_EQ(TF_GraphNumFunctions(func_graph_), 1);
  EXPECT_EQ(TF_GraphGetFunctions(func_graph_, funcs, 1, s_), 1);

// If `grad` is null and `func` is in `g`, TF_GraphCopyFunction is a noop.
// `grad` must have appropriate signature as described in the doc of
// GradientDef in tensorflow/core/framework/function.proto.
//
// If successful, status is set to OK and `func` and `grad` are added to `g`.
// Otherwise, status is set to the encountered error and `g` is unmodified.
TF_CAPI_EXPORT extern void TF_GraphCopyFunction(TF_Graph* g,