limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_
#define TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_

// A simple logging device to test custom device registration.
#include <memory>

#include "tensorflow/c/c_api.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_experimental.h"

  virtual ImmediateExecutionContext* GetContext() const = 0;

  // Following two methods are used to support custom device.
  // Return true if the inputs contain custom device tensor handle. It means
  // that the argument need to be handled by a custom device.
  virtual bool HasCustomDeviceInput() const = 0;

  virtual const tensorflow::OpDef* OpDef() const = 0;

//    void TF_InitGraph(TP_OptimizerRegistrationParams* params,
//                            TF_Status* status) {
//      params->struct_size = TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE;
//      params->device_type = "MY_DEVICE";
//
//      // Disable certain optimizer.
//      params->optimizer_configs->struct_size =
//      TP_OPTIMIZER_CONFIGS_STRUCT_SIZE; params->optimizer_configs->remapping =
//      TF_TriState_Off;
//

  // Experimental Custom Device.
  //===--------------------------------------------------------------------===//
  virtual CustomDeviceOpHandler& GetCustomDeviceOpHandler() = 0;

  // Returns whether `device_name` is registered as a custom device.
  virtual bool IsCustomDevice(const string& device_name) = 0;

  // Register a custom device. It will return error is the device name is
  // already registered.

TF_Operation* Mul(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
                  TF_Status* s, const char* name = "mul");

// If `op_device` is non-empty, set the created op on that device.
TF_Operation* MinWithDevice(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
                            const string& op_device, TF_Status* s,
                            const char* name = "min");

TF_Operation* Neg(TF_Operation* n, TF_Graph* graph, TF_Status* s,

// These APIs are mainly to facilitate testing and are subject to change.

// Records the op name in the `ForwardOperation`.
Status Reset(AbstractOperation*, const char* op, const char* raw_device_name,
             ForwardOperation*);

// Records the inputs in the `ForwardOperation`.
Status AddInput(AbstractOperation*, AbstractTensorHandle*, ForwardOperation*);
Status AddInputList(AbstractOperation*,