google / jaxopt

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

OSQP solvers with fun set and no init_params have misleading error message #545

Open jewillco opened 8 months ago

jewillco commented 8 months ago

The error message for setting fun to a function and not providing init_params to a class such as jaxopt.BoxOSQP is incorrect:

ValueError: init_x must be provided when fun is not None.

There is no init_x parameter to the run() method. Also, there are no examples in the documentation of those functions for what should be in a KKTSolution object, and that object appears to be undocumented.

jewillco commented 8 months ago

It would also be nice to be able to provide only the initial solution point but not the dual information. Also, which namespace the KKTSolution class is in does not appear to be listed either, and import jaxopt does not define it as either jaxopt.KKTSolution or jaxopt.base.KKTSolution.

Algue-Rythme commented 8 months ago

It would also be nice to be able to provide only the initial solution point but not the dual information

You can do this, as follows:

initial_solution_point = ... # your initial point
solver = BoxOSQP(...)
init_params = solver.init_params(init_x= initial_solution_point)  # if you don't want to initialize it with dual information.
solver.run(init_params=init_paramsl, ...)

I agree that more documentation is required regarding KKTSolution. @mblondel maybe add a generic documentation that explains some API choices?

jewillco commented 8 months ago

It also looks like you can't just fill in the initial point in KKTSolution. The primal field needs to have two entries, and it looks like the duals are required as well or else the implementation crashes. It looks like 0 might work in the other fields, but I am not sure. There also appears to be little shape checking in the implementation, so if a user passes parameters with mismatched sizes the failures are all internal to the library rather than pointing to user-defined arguments.

Algue-Rythme commented 8 months ago

The primal field needs to have two entries, and it looks like the duals are required as well or else the implementation crashes. It

That's true! The method init_params is there for this purpose: it fills the second field of primal variables and the dual variables. In theory you can warm start with any value for the primal since the problem is convex (the optimum does not depend on the initialization), I guess it is also true for dual variable but this needs to be checked.

There also appears to be little shape checking in the implementation, so if a user passes parameters with mismatched sizes the failures are all internal to the library rather than pointing to user-defined arguments.

You're right!

jewillco commented 8 months ago

Oh -- there's a method for filling in the parameters and it isn't just internal? Is that documented somewhere?

Algue-Rythme commented 8 months ago

Yes, it is part of the doc.

jewillco commented 8 months ago

I assumed that was an alternative to passing the parameters to run(), not a way to create them. Maybe there should be an example using run() with an init_params() call as an argument.

Algue-Rythme commented 8 months ago

For consistency reasons of the API we are forced to take params KKT tuple in run and update, we cannot support an additionnal init_x without breaking implicit differentiation. I'm afraid it won't change soon (if ever), so run signature will stay as is.

I agree we should add an example featuring init_params.