search

google / jaxopt

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

jittable verbose=True #503

Closed amir-saadat closed 1 year ago

amir-saadat commented 1 year ago

updated the print calls of the optimizers to jax.debug.print(..) so it can be jittable for closure definitions.

amir-saadat commented 1 year ago

thanks Mathieu, I've added a unit test that covers several solver, please let me know if I missed any. It would have been nice if we could automatically determine that the solver is in a jitted closure and we could automatically make jit=True. Let me know if you have any thoughts on that.

amir-saadat commented 1 year ago

a few things that I haven't tested here, any thoughts on how to test would be appreciated:

also asking this in case it was missed "It would have been nice if we could automatically determine that the solver is in a jitted closure and we could automatically make jit=True. Let me know if you have any thoughts on that." cause it might not be immediately obvious that the user needs to also pass jit=True or may sound unnecessary.

amir-saadat commented 1 year ago

printing in osqp

added a test for quadratic programming.

@mblondel please let me know if you want me to add anything else.