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);

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

  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);