Batch evaluation of problem data

[klamike]

The main idea for this right now is to use ExaModels to generate the kernels. However, we still want to keep the front-end API as JuMP models, so we need some way of converting JuMP functions (which are in terms of 1 instance) to ExaModels structs that take batches as input.

There already exists InfiniteExaModels.jl, which can be used to generate an ExaModel from an InfiniteOpt model, which ends up being pretty close to what we need -- the "number of supports" in InfiniteOpt is like our "batch size". However, I think it may be better to instead maintain our own single-instance JuMP function -> batched ExaModels.SIMDFunction pipeline. There are a few arguments for having our own pipeline:

• We'd have two fewer dependencies.
• If we use InfiniteExaModels, we either need to implement JuMP -> InfiniteOpt or force users to give an InfiniteOpt model as input. This may not be as compatible with existing formulation libraries, but (using a subset of) their API is actually quite nice for our purpose.
• It allows us greater flexibility in the structure of the generated ExaModels.SIMDFunctions, as well as their evaluation. That can be important for e.g. jacobians since we'd get huge block-diagonal ones using InfiniteExaModels since ExaModels doesn't realize the batch instances are totally independent (in fact they are not, since they'd be tied together using an expectation in the objective). Similar story for the hessian of the Lagrangian. Indexing can also be made nicer; under the InfiniteExaModels setup, we have to flatten the predictions/reshape the results.
• Doing the conversion at the function level rather than the model level more closely aligns with expected use cases in L2O. For example, the completion of the special cases in L2ODLL evaluates JuMP functions that are built during the decomposition. It is not obvious how we'd use InfiniteExaModels to generate kernels for those, but it is easy if we have a JuMP -> SIMDFunction pipeline.

I am prototyping some of this at https://github.com/klamike/BatchOptInterface.jl

cc @frapac @andrewrosemberg

[klamike]

https://github.com/LearningToOptimize/L2ODLL.jl/blob/bb74ad02bb6b598da0608293a238db4eb5a2d3ca/src/batch.jl has some special cases for batched AffExpr, QuadExpr, and Vector{AffExpr} via broadcasting.
Once exanauts/ExaModels.jl#148 is merged, I'll add batched NonlinearExpr and Vector{NonlinearExpr}

[klamike]

https://github.com/klamike/BatchNLPKernels.jl is looking pretty good now! I want to eventually move the code in L2ODLL linked above to something like BatchConicKernels.jl, for e.g. distance to set and specific kernels for batched Ax-b x'Qx + c'x etc.

May eventually combine them into e.g. L2OKernels.jl and put them in here under lib/, but I think it's better to implement end-to-end a few L2O methods so we can get a sense of what the final structure should look like for this repo, to avoid repeatedly refactoring.