External solvers in JAXOPT

[SNMS95]

Hi,

I have to use an external solver (from PETSc) to solve a linear system in the forward pass.
In order to calculate the gradients, can I use jaxopt.implicit_diff.root_jvp decorator for this purpose?

I could not find a relevant documentation. If there is one, please let me know.
Thanks in advance

[mblondel]

Here's an example on how to do it in the case of unconstrained optimization:

https://jaxopt.github.io/stable/auto_examples/implicit_diff/ridge_reg_implicit_diff.html#sphx-glr-auto-examples-implicit-diff-ridge-reg-implicit-diff-py

[SNMS95]

Hi @mblondel , Thanks for the answer.

I looked at the example. However, the solve parameter of jaxopt.implicit_diff.custom_root is not set here. Is the solve parameter used to calculate the adjoints for autodiff?

If so, Does the solve has to be a JAXOPT solver ?

[mblondel]

Do you want to replace the linear solver in the forward (i.e. your function outputs the solution of a linear system) or in the backward pass (i.e. you want to change the solver used for implicit diff)?

[SNMS95]

At the very least, I want to replace the one in the forward pass.
But if possible, I would like to replace both of them.

[mblondel]

The forward solver is the one that you decorate with custom_root (ridge_solver in the example above).
The backward solver is the one you pass as the solve option. It should have the form solve(matvec, b) where matvec is the linear operator you want to use.

[SNMS95]

Thanks :)
It is clear now!