Differentiation of quadratic problem with parameters tuning quadratic coefficient using Gurobi and quadratic interface

[asharamanujam]

DiffOpt seems to have a problem with using the quadratic interface to tune quadratic coefficients in a quadratic problem using the quadratic interface. If I use the following:

using JuMP
using Gurobi
using DiffOpt, Zygote, Flux
using MathOptInterface


const MOI = MathOptInterface


using ChainRulesCore
function tunable_layer_quadratic(alpha,beta;
    model::Model = DiffOpt.quadratic_diff_model(Gurobi.Optimizer))
    @variable(model, alpha_var[i in 1:4] in Parameter(alpha[i]))
    @variable(model, beta_var[i in 1:4] in Parameter(beta[i]))
    @variable(model,0<=x[1:4]<=5)
    @constraint(model,x[1]+x[2]<=3)
    @constraint(model,x[2]+x[3]+x[4]>=2)
    @constraint(model,x[4]+x[1]<=2.5)
    @objective(model,Min,sum(alpha_var[i]*x[i]*x[i] + beta_var[i]*x[i] for i in 1:4))
    optimize!(model)
    return value.(model[:x])
end


function ChainRulesCore.rrule(::typeof(tunable_layer_quadratic),alpha,beta)
    opt = MOI.OptimizerWithAttributes(
        Gurobi.Optimizer,
        "QCPDual" => 1,
    )
    model = DiffOpt.quadratic_diff_model(opt)
    output = tunable_layer_quadratic(alpha,beta;model = model)
    alpha_var = model[:alpha_var]
    beta_var = model[:beta_var]
    alpha_param_refs = [ParameterRef(alpha_var[i]) for i in 1:4]
    beta_param_refs = [ParameterRef(beta_var[i]) for i in 1:4]
    function pullback_tunable_layer(dx)
        x_var = model[:x]
        println("dx: ",dx)
        MOI.set.(model, DiffOpt.ReverseVariablePrimal(), x_var, dx)


        DiffOpt.reverse_differentiate!(model)
        d_alpha = MOI.get.(model, DiffOpt.ReverseConstraintSet(), alpha_param_refs)
        d_beta = MOI.get.(model, DiffOpt.ReverseConstraintSet(), beta_param_refs)


        grad_dalpha = [g.value for g in d_alpha]
        grad_dbeta = [g.value for g in d_beta]


        return NoTangent(), grad_dalpha,grad_dbeta
    end
    return output, pullback_tunable_layer
end


function loss_fn(alpha,beta)
    output = tunable_layer_quadratic(alpha,beta)
    return sum(output)
end
alpha = [1.0, 2.0, 3.0, 4.0]
beta = [0.5, 1.0, 1.5, 2.0]


ps = Flux.params(alpha, beta)


loss, back = Flux.pullback(() -> loss_fn(alpha, beta), ps)
gs = back(1.0)


println("loss = ", loss)
println("∂loss/∂alpha = ", gs[alpha])
println("∂loss/∂beta  = ", gs[beta])

I get the following error: ERROR: LoadError: The solver does not support an objective function of type MathOptInterface.ScalarNonlinearFunction.

If I modify the code and use the following:

function tunable_layer_quadratic_1(alpha,beta;
    model::Model = DiffOpt.quadratic_diff_model(Gurobi.Optimizer))
    @variable(model, alpha_var[i in 1:4] in Parameter(alpha[i]))
    @variable(model, beta_var[i in 1:4] in Parameter(beta[i]))
    @variable(model,0<=x[1:4]<=5)
    @constraint(model,x[1]+x[2]<=3)
    @constraint(model,x[2]+x[3]+x[4]>=2)
    @constraint(model,x[4]+x[1]<=2.5)
    @variable(model,t_x[1:4]>=0)
    for i in 1:4
        @constraint(model,t_x[i]>=x[i]^2)
    end
   
    @objective(model,Min,sum(alpha_var[i]*t_x[i] + beta_var[i]*x[i] for i in 1:4))
    optimize!(model)
    return value.(model[:x])
end


function ChainRulesCore.rrule(::typeof(tunable_layer_quadratic_1),alpha,beta)
    opt = MOI.OptimizerWithAttributes(
        Gurobi.Optimizer,
        "QCPDual" => 1,
    )
    model = DiffOpt.quadratic_diff_model(opt)
    output = tunable_layer_quadratic_1(alpha,beta;model = model)
    alpha_var = model[:alpha_var]
    beta_var = model[:beta_var]
    alpha_param_refs = [ParameterRef(alpha_var[i]) for i in 1:4]
    beta_param_refs = [ParameterRef(beta_var[i]) for i in 1:4]
    function pullback_tunable_layer(dx)
        x_var = model[:x]
        println("dx: ",dx)
        MOI.set.(model, DiffOpt.ReverseVariablePrimal(), x_var, dx)


        DiffOpt.reverse_differentiate!(model)
        d_alpha = MOI.get.(model, DiffOpt.ReverseConstraintSet(), alpha_param_refs)
        d_beta = MOI.get.(model, DiffOpt.ReverseConstraintSet(), beta_param_refs)


        grad_dalpha = [g.value for g in d_alpha]
        grad_dbeta = [g.value for g in d_beta]


        return NoTangent(), grad_dalpha,grad_dbeta
    end
    return output, pullback_tunable_layer
end


function loss_fn_1(alpha,beta)
    output = tunable_layer_quadratic_1(alpha,beta)
    return sum(output)
end


alpha = [1.0, 2.0, 3.0, 4.0]
beta = [0.5, 1.0, 1.5, 2.0]


ps = Flux.params(alpha, beta)


loss, back = Flux.pullback(() -> loss_fn_1(alpha, beta), ps)
gs = back(1.0)


println("loss = ", loss)
println("∂loss/∂alpha = ", gs[alpha])
println("∂loss/∂beta  = ", gs[beta])

I get the following error: ERROR: LoadError: UnsupportedConstraint: MathOptInterface.ScalarQuadraticFunction{Float64}-in-MathOptInterface.GreaterThan{Float64} constraints are not supported by the
solver you have chosen, and we could not reformulate your model into a
form that is supported.

To fix this error you must choose a different solver.

Is it possible to use the quadratic interface and Gurobi to find sensitivities with respect to quadratic coefficients? If so, can you help me out?