CSession::~CSession() {
  TF_Status* s = TF_NewStatus();
  CloseAndDelete(s);
  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
  TF_DeleteStatus(s);
}

void CSession::SetInputs(
    std::vector<std::pair<TF_Operation*, TF_Tensor*>> inputs) {
  DeleteInputValues();
  inputs_.clear();
  for (const auto& p : inputs) {
    inputs_.emplace_back(TF_Output{p.first, 0});

class CSession {
 public:
  CSession(TF_Graph* graph, TF_Status* s, bool use_XLA = false);
  explicit CSession(TF_Session* session);

  ~CSession();

  void SetInputs(std::vector<std::pair<TF_Operation*, TF_Tensor*>> inputs);
  void SetOutputs(std::initializer_list<TF_Operation*> outputs);
  void SetOutputs(const std::vector<TF_Output>& outputs);

  // Create a session for this graph.
  CSession csession(graph, s);
  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);

  // Run the graph.
  csession.SetInputs({{feed, Int32Tensor(3)}});
  csession.SetOutputs({add});
  csession.Run(s);
  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
  TF_Tensor* out = csession.output_tensor(0);
  ASSERT_TRUE(out != nullptr);

    // Create a session for this graph
    CSession csession(host_graph_, s_);
    ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);

    // Run
    csession.SetInputs(inputs);
    csession.SetOutputs(outputs);
    csession.Run(s_);
    ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);

    // Check results
    for (int i = 0; i < expected_results.size(); ++i) {

class ImmediateExecutionOperation : public AbstractOperation {
 public:
  virtual void Clear() = 0;

  // Returns the inputs of this op.
  virtual absl::Span<ImmediateExecutionTensorHandle* const> GetInputs()
      const = 0;
  virtual Status SetInput(size_t index,
                          ImmediateExecutionTensorHandle* input) = 0;

  virtual ImmediateExecutionContext* GetContext() const = 0;

}

extern int TFE_OpGetFlatInputCount(const TFE_Op* op, TF_Status* status) {
  return tensorflow::unwrap(op)->GetInputs().size();
}

extern TFE_TensorHandle* TFE_OpGetFlatInput(const TFE_Op* op, int index,
                                            TF_Status* status) {
  return tensorflow::wrap(tensorflow::unwrap(op)->GetInputs()[index]);
}

TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name,