TF_CAPI_EXPORT extern TF_Buffer* TF_GetRegisteredKernelsForOp(
    const char* name, TF_Status* status);

// Update edge, switch input/ output in a node
TF_CAPI_EXPORT extern void TF_UpdateEdge(TF_Graph* graph, TF_Output new_src,
                                         TF_Input dst, TF_Status* status);

// --------------------------------------------------------------------------

  status->status = MessageToBuffer(kernel_list, ret);
  if (!status->status.ok()) {
    TF_DeleteBuffer(ret);
    return nullptr;
  }
  return ret;
}

void TF_UpdateEdge(TF_Graph* graph, TF_Output new_src, TF_Input dst,
                   TF_Status* status) {
  using tensorflow::RecordMutation;
  mutex_lock l(graph->mu);
  tensorflow::shape_inference::InferenceContext* ic =

  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);

  NodeDef node_def_neg;
  ASSERT_TRUE(GetNodeDef(neg, &node_def_neg));
  EXPECT_EQ(string("add"), node_def_neg.input(0));

  // update edge of neg
  TF_UpdateEdge(graph, TF_Output{one, 0}, TF_Input{neg, 0}, s);

  ASSERT_TRUE(GetNodeDef(neg, &node_def_neg));
  EXPECT_EQ(string("one:0"), node_def_neg.input(0));

  // Clean up
  TF_DeleteGraph(graph);