}

  // Populate return_nodes
  DCHECK(tf_results->return_nodes.empty());
  tf_results->return_nodes.resize(results.return_nodes.size());
  for (int i = 0; i < results.return_nodes.size(); ++i) {
    tf_results->return_nodes[i] = ToOperation(results.return_nodes[i]);
  }

  // Populate missing unused map keys
  DCHECK(tf_results->missing_unused_key_names.empty());
  DCHECK(tf_results->missing_unused_key_indexes.empty());

  TFE_ContextOptions* opts = TFE_NewContextOptions();
  TFE_Context* ctx = TFE_NewContext(opts, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_DeleteContextOptions(opts);

  TFE_TensorHandle* h = TestScalarTensorHandle(ctx, 1.0f);
  TFE_TensorDebugInfo* debug_info = TFE_TensorHandleTensorDebugInfo(h, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

}

TF_AbstractTensor* TF_AddFunctionParameter(TF_ExecutionContext* func,
                                           TF_DataType dtype, TF_Shape shape,
                                           TF_Status* s) {
  DCHECK_GE(shape.num_dims, -1);
  TracingTensorHandle* t;
  TracingContext* tracing_ctx = dyn_cast<TracingContext>(unwrap(func));
  if (!tracing_ctx) {
    tsl::Set_TF_Status_from_Status(

     *
     *    dA = U
     *    dbias = reduceSum(U, dims = channel_dim)
     */

    AbstractTensorHandle* upstream_grad = grad_outputs[0];
    DCHECK(upstream_grad);

    // Recover data format from forward pass for gradient.
    std::string data_format;
    TF_RETURN_IF_ERROR(forward_attrs_.Get("data_format", &data_format));

    // Grad for A

    t.reset(x_raw);
  }

  AbstractOperationPtr check_numerics_op(ctx->CreateOperation());
  ForwardOperation forward_op;
  Status s = Reset(check_numerics_op.get(), "CheckNumerics",
                   /*raw_device_name=*/nullptr, &forward_op);
  ASSERT_EQ(errors::OK, s.code()) << s.message();
  if (isa<TracingOperation>(check_numerics_op.get())) {
    s = dyn_cast<TracingOperation>(check_numerics_op.get())
            ->SetOpName("check_numerics");

using ::testing::NotNull;
using ::tsl::testing::StatusIs;

PJRT_Buffer* CreateCBuffer() {
  auto status = pjrt::PjrtApi(DEVICE_CPU);
  if (!status.ok()) {
    CHECK_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
  }
  auto pjrt_client = xla::GetCApiClient(DEVICE_CPU);
  CHECK_OK(pjrt_client.status());
  auto c_api_client = down_cast<xla::PjRtCApiClient*>(pjrt_client->get());
  std::vector<int32_t> data(1, 0);

  TF_ASSERT_OK(cluster->Provision());

  TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
                                          cluster->GetDeviceNames());
  GrapplerItem item;
  CHECK(fake_input.NextItem(&item));

  TF_Status* status = TF_NewStatus();
  TF_GraphProperties* graph_properties =
      TF_NewGraphProperties(reinterpret_cast<TF_GrapplerItem*>(&item));

namespace {

using ::tensorflow::string;

// Add the values of three variables on three different tasks.
string AddVariablesFunction() {
  tensorflow::FunctionDef def;
  CHECK(tensorflow::protobuf::TextFormat::ParseFromString(
      "    signature {"
      "      name: 'AddVariablesFunction'"
      "      input_arg {"
      "        name: 'var'"
      "        type: DT_RESOURCE"
      "      }"

  // Initialize auxilary data structures.
  vector<float> thetaPlus_data(num_elems);
  vector<float> thetaMinus_data(num_elems);
  AbstractTensorHandle* f_outputs[1];

  // Numerical Grad Check
  for (int i = 0; i < num_elems; i++) {
    // Get relative epsilon value
    float epsilon = theta_data[i] == 0 ? 1e-4 : std::abs(theta_data[i] * 1e-4);
    AbstractTensorHandlePtr two_eps;
    {

  control_output_nodes.reserve(ncontrol_outputs);
  for (int i = 0; i < ncontrol_outputs; ++i) {
    control_output_nodes.push_back(&control_outputs[i]->node);
  }

  // Do the actual function creation.
  DCHECK(append_hash_to_fn_name <= 1);
  tensorflow::FunctionDef fdef;
  status->status = tensorflow::GraphToFunctionDef(
      fn_body->graph, fn_name, append_hash_to_fn_name != 0,
      /*set_stateful_from_nodes=*/true,