Batched QP (and other optimization algorithm)

[jn-tang]

I'm trying to make OSQP batchable (so I can make it a layer in neural networks, like OptNet), but I couldn't find any documentation yet about using vmap to solve batched version of optimization problems.

[Algue-Rythme]

Hi !

OSQP is batchable if you set the jitargument in constructor to True. The functions init_state, init_params, update and run are pure and jittable. Hence you can wrap them in any other function that you want to jit/vmap.

However note that a single OSQP instance corresponds to a certain form of problem with fixed hyper-parameters (in constructor). Hence all the instances you want to solve must share the same hyper-parameters. Only the initialization and the params_obj,params_eq,params_ineq may vary from an instance to another.