Add Batched{S} set

[odow]

@amontoison wants this. There are also some folks in the GPU/cvxpy domain sniffing around this (cvxpy/cvxpy#2485).

We have almost everything we need. There's just a choice between two approaches.

First, we could set the value of Parameter to a vector of sets:

MOI.set(
    ::Optimizer,
    ::MOI.ConstraintSet,
    ::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}},
    ::Vector{MOI.Parameter{T}},
)

But this might be hard to get through the various MOI layers.

The easier alternative is to add a new Batched{S<:MOI.AbstractSet} set.

If there are multiple batched sets, they all must have the same length. Then the results are returned via result_count and we don't support returning multiple solutions within a batch.

Then a solver could choose to natively support batched parameters. And we could have a fallback optimizer like:

using JuMP
import Ipopt
import MathOptBatchOptimizer as MOBO
model = Model(() -> MOBO.Optimizer(Ipopt.Optimizer))
@variable(model, x)
@variable(model, p in Parameter(1))
set_parameter_value(p, 1:100) # <-- set as a vector, not a scalar
@objective(model, Min, x)
@constraint(model, x >= p)
optimize!(model)
@assert result_count(model) == 100

[aman-coder03]

hi @odow i would like to work on this
i am leaning toward adding Batched{S<:MOI.AbstractSet} since it seems cleaner than pushing vectors of sets through MOI

[blegat]

I like the idea of the Batched set. It should work quite well with result_count as well. POI could do a for loop in case the solver does not supports the set Batched{S<:MOI.AbstractSet}

[aman-coder03]

@blegat i will proceed with introducing a new Batched{S<:MOI.AbstractSet} set

my plan is:
define struct Batched{S<:MOI.AbstractSet} <: MOI.AbstractSet that wraps an underlying set S and stores a batch_size::Int
enforce that all Batched sets added to a model share the same batch_size (validated at the optimizer level)
expose batched results through result_count, with each batch corresponding to one result index
do not support multiple solutions within a batch, as mentioned in the issue
leave room for solvers to natively implement support for Batched, while a fallback optimizer (e.g., a POI-style wrapper) can loop over the batch dimension if unsupported

before I start implementing should the batch_size live explicitly inside the Batched set (as a field), or would you prefer that it be inferred from the associated parameter values?

[blegat]

If Batched contains a vector of sets then, the batch_size would be the length of the vector, I'd prefer not store batch_size as well

[aman-coder03]

i will define Batched{S<:MOI.AbstractSet} to store a vector of sets and infer the batch size from length(sets) without storing it separately
i will enforce consistency of the batch length across all Batched sets at the optimizer level and wire result_count to match that length
started working on the initial implementation and soon i will open a PR

[odow]

Another reason to add this: Gurobi has support for multiple scenarios: https://docs.gurobi.com/projects/optimizer/en/current/features/multiscenario.html

It might also be useful for SDDP.jl.

[blegat]

We should also implement the Batched set in https://github.com/JuliaSmoothOptimizers/NLPModelsJuMP.jl, which should then translate it to the NLPModels' equivalent (JuliaSmoothOptimizers/NLPModels.jl#521), which would then for example be used by MadNLP (MadNLP/MadIPM.jl#78)

[blegat]

We added a GSOC proposal on this: https://github.com/jump-dev/GSOC

[joaquimg]

I like the Batched set idea.

How would it appear in the JuMP level?

I think this is too mystical to just work:
@variable(model, p in Parameter(1))
Maybe this is more explicit:
@variable(model, p in Batched{Parameter}(1))
It would just foward the inputs of the call to the inner set.
I dont we need the batch size a priori, it could be flexible.

I also think this is error prone:
set_parameter_value(p, 1:100) # <-- set as a vector, not a scalar
can we be explicit? for instance:
set_batched_parameter_value(p, 1:100)