TF_RETURN_IF_ERROR(
        RunAndMaybeSum(ctx, forward, theta_inputs, f_outputs, use_function));
    AbstractTensorHandlePtr fPlus(f_outputs[0]);

    // Get f(theta - eps):
    theta_inputs[input_index] = thetaMinus.get();
    TF_RETURN_IF_ERROR(
        RunAndMaybeSum(ctx, forward, theta_inputs, f_outputs, use_function));
    AbstractTensorHandlePtr fMinus(f_outputs[0]);

                                    TF_Graph* graph, TF_Operation* ctrl_op,
                                    TF_Status* s, const char* name = "add");

TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s,
                  const char* name = "add");

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

  TF_SetRequestedDevice(graph, op, device);
}

void UpdateEdge(TF_Graph* graph, TF_Output new_src, TF_Input dst,
                TF_Status* status) {
  TF_UpdateEdge(graph, new_src, dst, status);
}

void ExtendSession(TF_Session* session, TF_Status* status) {
  TF_ExtendSession(session, status);
}

std::string GetHandleShapeAndType(TF_Graph* graph, TF_Output output) {
  Node* node = &output.oper->node;

std::string GetHandleShapeAndType(TF_Graph* graph, TF_Output output);

// Sets `output` based on `proto`, which should be a serialized
// CppShapeInferenceResult::HandleData proto. `output` should be a resource
// or variant tensor.
// NOTE(skyewm): `proto` is passed a void*/size_t pair instead of a std::string
// because I couldn't get SWIG to work otherwise.
void SetHandleShapeAndType(TF_Graph* graph, TF_Output output, const void* proto,

namespace tensorflow {
namespace gradients {
namespace {

class ReluGradientFunction : public GradientFunction {
 public:
  explicit ReluGradientFunction(vector<AbstractTensorHandle*> f_outputs)
      : forward_outputs_(f_outputs) {
    for (auto output : forward_outputs_) {
      if (output) {
        output->Ref();
      }
    }
  }

  Status Compute(AbstractContext* ctx,

  // Used to link graphs contained in TF_WhileParams to the parent graph that
  // will eventually contain the full while loop.
  TF_Graph* parent;
  TF_Output* parent_inputs;
};

struct TF_OperationDescription {
  TF_OperationDescription(TF_Graph* g, const char* op_type,
                          const char* node_name)

    std::vector<AbstractTensorHandle*> fn_outputs(retvals);
    TF_RETURN_IF_ERROR(fn_op->Execute(
        absl::Span<AbstractTensorHandle*>(fn_outputs.data(), fn_outputs.size()),
        &retvals));
    int skipped_indices = 0;
    for (int i = 0; i < outputs.size(); i++) {
      if (!null_indices.contains(i)) {
        outputs[i] = fn_outputs[i - skipped_indices];
      } else {
        skipped_indices += 1;

namespace tensorflow {

REGISTER_OP("MlirPassthroughOp")
    .Attr("mlir_module: string")
    .Attr("Tinputs : list(type) >= 0")
    .Input("inputs: Tinputs")
    .Attr("Toutputs : list(type) >= 0")
    .Output("outputs: Toutputs")
    .SetShapeFn(shape_inference::UnknownShape);

// Check that attributes that define types are exported.

// CHECK: key: "Tinputs"
// CHECK-NEXT:    value
// CHECK-NEXT:      list
// CHECK-NEXT:        type: DT_FLOAT

// CHECK: key: "Toutputs"
// CHECK-NEXT:    value
// CHECK-NEXT:      list
// CHECK-NEXT:        type: DT_FLOAT

// CHECK: "extra_type_attr"
// CHECK-NEXT:    value
// CHECK-NEXT:      list
// CHECK-NEXT:        type: DT_INT32

    }
    if (TF_GetCode(s) != TF_OK) return;
  }
  std::vector<TFE_TensorHandle*> op_outputs(*num_outputs);
  TFE_Execute(op, op_outputs.data(), num_outputs, s);
  TFE_DeleteOp(op);
  if (TF_GetCode(s) != TF_OK) return;
  std::vector<TFE_TensorHandle*> unwrapped_outputs;
  unwrapped_outputs.reserve(op_outputs.size());
  for (auto* handle : op_outputs) {
    unwrapped_outputs.push_back(handle);
  }