TFE_NewContextOptions(), TFE_DeleteContextOptions);
  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
      TF_CreateConfig(
          /*enable_xla_compilation=*/false,
          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
      TF_DeleteBuffer);
  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
                              status.get());

unsigned char TF_SetXlaEnableLazyCompilation(unsigned char enable) {
  tensorflow::BuildXlaOpsPassFlags* flags =
      tensorflow::GetBuildXlaOpsPassFlags();
  bool original = flags->tf_xla_enable_lazy_compilation;
  flags->tf_xla_enable_lazy_compilation = enable;
  return original;
}

unsigned char TF_SetTfXlaCpuGlobalJit(unsigned char enable) {
  tensorflow::MarkForCompilationPassFlags* flags =
      tensorflow::GetMarkForCompilationPassFlags();

      TFE_NewContextOptions(), TFE_DeleteContextOptions);
  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
      TF_CreateConfig(
          /*enable_xla_compilation=*/false,
          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
      TF_DeleteBuffer);
  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
                              status.get());

}

// All the *Helper methods are used as a workaround for the restrictions that
// one cannot call ASSERT_* methods in non-void-returning functions (when
// exceptions are disabled during compilation)
void PlaceholderHelper(TF_Graph* graph, TF_Status* s, const char* name,
                       TF_DataType dtype, const std::vector<int64_t>& dims,
                       TF_Operation** op) {

  TF_ExecutionContextGetTFEContext(graph_ctx, status.get());
  ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(status.get()));

  TF_DeleteExecutionContext(graph_ctx);
}

// The above tests are run for a combination of:
// - graphdef and MLIR tracing engine
INSTANTIATE_TEST_SUITE_P(Tracing, UnifiedCAPI,
                         ::testing::Combine(::testing::Values("graphdef",

  TF_Operation* plusB = Add(plus2, b, graph, s, "plusB");
  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);

  // Setup a session and a partial run handle.  The partial run will allow
  // computation of A + 2 + B in two phases (calls to TF_SessionPRun):
  // 1. Feed A and get (A+2)
  // 2. Feed B and get (A+2)+B
  TF_SessionOptions* opts = TF_NewSessionOptions();
  TF_Session* sess = TF_NewSession(graph, opts, s);

    return op;
  }

  // All the *Helper methods are used as a workaround for the restrictions that
  // one cannot call ASSERT_* methods in non-void-returning functions (when
  // exceptions are disabled during compilation)
  void UseHelper(const std::vector<TF_Output>& inputs, TF_Operation** op) {
    TF_OperationDescription* desc =
        TF_NewOperation(host_graph_, func_name_, func_node_name_);
    for (auto input : inputs) {