search

google / jaxopt

Hardware accelerated, batchable and differentiable optimizers in JAX.
https://jaxopt.github.io
Apache License 2.0
930 stars 65 forks source link

Batched QP (and other optimization algorithm) #347

Open jn-tang opened 1 year ago

jn-tang commented 1 year ago

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 commented 1 year ago

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.