// These XLA flags are needed to trigger XLA properly from C (more generally
    // non-Python) clients. If this API is called again with `enable` set to
    // false, it is safe to keep these flag values as is.
    tensorflow::MarkForCompilationPassFlags* flags =
        tensorflow::GetMarkForCompilationPassFlags();
    flags->tf_xla_cpu_global_jit = true;
    flags->tf_xla_min_cluster_size = 1;
  } else {

    sum_dims.reset(sum_dims_raw);
  }

  // Reduce sum the output on all dimensions.
  TF_RETURN_IF_ERROR(ops::Sum(ctx, model_out.get(), sum_dims.get(), &outputs[0],
                              /*keep_dims=*/false, "sum_output"));
  return absl::OkStatus();
}
// ========================= End Helper Functions==============================

Status CalcNumericalGrad(AbstractContext* ctx, Model forward,

  //   Outputs:
  std::vector<TensorHandlePtr> op_outputs_ TF_GUARDED_BY(execution_mutex_);
  // TF_Status is an incomplete type and so can't be stack allocated. To avoid
  // unnecessary allocations each Execute call, we keep one heap-allocated
  // version for the thread.
  StatusPtr status_ TF_GUARDED_BY(execution_mutex_);

  const std::string device_;
  ExecutorPtr executor_ TF_GUARDED_BY(execution_mutex_);

                                                   int keep_alive_secs,
                                                   const void* proto,
                                                   size_t proto_len,
                                                   TF_Status* status) {
  TFE_ContextSetServerDefWithTimeoutAndRetries(
      ctx, keep_alive_secs, proto, proto_len, /*init_timeout_in_ms=*/0,

  if (use_function) {
    const char* fn_name = "test_fn";
    core::RefCountPtr<AbstractFunction> scoped_func;
    // Returning null tensors from a tf.function is not supported, so we keep
    // track of indices in the model's outputs are nullptr in this set.
    // The FunctionDef only outputs the non-null tensors. We later pad the
    // function op outputs to have nullptrs at the `null_indices`.

    return;
  }
  status->status = coord_agent->WaitAtBarrier(
      barrier_id, absl::Milliseconds(barrier_timeout_in_ms), {});
}

void TFE_InitializeLocalOnlyContext(TFE_Context* ctx, int keep_alive_secs,
                                    const void* proto, size_t proto_len,
                                    TF_Status* status) {
  tensorflow::ServerDef server_def;

  TFE_OpAddInput(op, input, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_OpAddInput(op, axis, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_OpSetAttrBool(op, "keep_dims", 1);
  TFE_OpSetAttrType(op, "Tidx", TF_INT32);
  TF_DeleteStatus(status);
  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(input));

  return op;
}