bool use_function) {
  AbstractTensorHandle* model_outputs[1];

  // Run the model.
  TF_RETURN_IF_ERROR(
      RunModel(forward, ctx, inputs, model_outputs, use_function));
  AbstractTensorHandlePtr model_out(model_outputs[0]);

  TF_Tensor* model_out_tensor;
  TF_RETURN_IF_ERROR(GetValue(model_out.get(), &model_out_tensor));
  int num_dims_out = TF_NumDims(model_out_tensor);

    // Get transpose attrs
    bool t_a;
    TF_RETURN_IF_ERROR(forward_attrs_.Get("transpose_a", &t_a));

    bool t_b;
    TF_RETURN_IF_ERROR(forward_attrs_.Get("transpose_b", &t_b));

    // Conj each input
    AbstractTensorHandle* conj_output;
    std::string name = "Conj_A_MatMul_Grad";
    TF_RETURN_IF_ERROR(
        SafeConj(ctx, forward_inputs_[0], &conj_output, name.c_str()));

    TF_RETURN_IF_ERROR(fn_op->Reset(fn_name, /*raw_device_name=*/nullptr));
    for (auto input : inputs) {
      TF_RETURN_IF_ERROR(fn_op->AddInput(input));
    }
    int retvals = outputs.size() - null_indices.size();
    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;

  ImmediateTensorHandlePtr dim(imm_ctx->CreateLocalHandle(minus_1.get()));
  AbstractTensorHandle* expand_dims_outputs;
  TF_RETURN_IF_ERROR(
      ops::ExpandDims(ctx, vec, dim.get(), &expand_dims_outputs, "ExpandDims"));
  TF_RETURN_IF_ERROR(
      ops::Mul(ctx, expand_dims_outputs, mat, &outputs[0], "Mul"));
  expand_dims_outputs->Unref();
  return absl::OkStatus();
}

  AbstractOperationPtr op(ctx->CreateOperation());
  TF_RETURN_IF_ERROR(op->Reset("ZerosLike", /*raw_device_name=*/nullptr));
  if (isa<tracing::TracingOperation>(op.get())) {
    TF_RETURN_IF_ERROR(dyn_cast<tracing::TracingOperation>(op.get())->SetOpName(
        absl::StrCat("ZerosLike", ToId(t)).c_str()));
  }
  TF_RETURN_IF_ERROR(op->AddInput(t));
  int num_outputs = 1;

    TracingOperationPtr operation(CreateOperation());
    TF_RETURN_IF_ERROR(operation->Reset("Placeholder", nullptr));
    TF_RETURN_IF_ERROR(
        operation->SetOpName(absl::StrCat("_input_", inputs_.size()).c_str()));
    TF_RETURN_IF_ERROR(operation->SetAttrType("dtype", dtype));
    if (!shape.unknown_rank()) {
      TF_RETURN_IF_ERROR(operation->SetAttrShape(

namespace tensorflow {
namespace tracing {

Status MaybeSetOpName(AbstractOperation* op, const char* op_name) {
  if (isa<TracingOperation>(op)) {
    TF_RETURN_IF_ERROR(dyn_cast<TracingOperation>(op)->SetOpName(op_name));
  }
  if (isa<gradients::TapeOperation>(op)) {
    TF_RETURN_IF_ERROR(MaybeSetOpName(
        dyn_cast<gradients::TapeOperation>(op)->GetBackingOperation(),
        op_name));
  }
  return absl::OkStatus();
}

  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);
  TFE_Context* eager_ctx =
      TF_ExecutionContextGetTFEContext(wrap(ctx), status.get());
  TF_RETURN_IF_ERROR(StatusFromTF_Status(status.get()));
  TFE_TensorHandle* input_eager =
      TestTensorHandleWithDims<T, datatype>(eager_ctx, data, dims, num_dims);
  *tensor =

  tape.Watch(inputs[0]);  // Watch x.
  AbstractTensorHandle* exp_output;
  TF_RETURN_IF_ERROR(ops::Exp(ctx, inputs[0], &exp_output, "Exp"));
  std::unique_ptr<GradientFunction> gradient_function(
      new PassThroughGradientFunction);
  tape.RecordOperation(inputs, {exp_output}, gradient_function.release());
  TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx,
                                          /*targets*/ {exp_output},

    }
    std::vector<AbstractTensorHandle*> temp_outputs(1);
    AbstractContextPtr tape_ctx(new TapeContext(ctx, &tape, grad_registry));
    TF_RETURN_IF_ERROR(
        forward_model(tape_ctx.get(), inputs, absl::MakeSpan(temp_outputs)));

    TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx, /*targets=*/temp_outputs,
                                            /*sources=*/inputs,