absl::Span<AbstractTensorHandle*> grad_inputs) override {
    AbstractTensorHandle* upstream_grad = grad_outputs[0];
    AbstractTensorHandle* activations = forward_outputs_[0];

    // Calculate Grad
    std::string name = "relu_grad";
    TF_RETURN_IF_ERROR(ReluGrad(ctx, upstream_grad, activations,
                                &grad_inputs[0], name.c_str()));
    return absl::OkStatus();
  }

    TF_RETURN_IF_ERROR(
        ops::Sub(ctx, fPlus.get(), fMinus.get(), f_outputs, "sub_top"));
    AbstractTensorHandlePtr fDiff(f_outputs[0]);

    // Calculate using the difference quotient definition:
    // (f(theta + eps) - f(theta - eps)) / (2 * eps).
    TF_RETURN_IF_ERROR(
        ops::Div(ctx, fDiff.get(), two_eps.get(), f_outputs, "diff_quotient"));

    AbstractTensorHandlePtr Ones_X(temp_output);

    name = "Add_Log1p_Grad_X";
    // Calculate 1 + Conj(X)
    TF_RETURN_IF_ERROR(
        AddV2(ctx, Ones_X.get(), Conj_X.get(), &temp_output, name.c_str()));

    AbstractTensorHandlePtr Conj_XP1(temp_output);

    name = "Div_Log1p_Grad_X";
    // Calculate U / (1 + Conj(X))
    TF_RETURN_IF_ERROR(