Problem differentiating through `solver.run` in `OptaxSolver`
- Dominant language
- Python
- Stars
- 1.1k
- Forks
- 76
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 1
Description
I've been trying to use `OptaxSolver` to perform a simple function minimization, since I want to differentiate through it's solution (the fixed point of the solver), but ran into an issue I'm not familiar with.
Here's a MWE for the error message:
```python3
import jax
import jax.scipy as jsp
from jaxopt import OptaxSolver
import optax
def pipeline(param_for_grad, data):
def to_minimize(latent):
return -jsp.stats.norm.logpdf(data, loc=param_for_grad*latent, scale=1)
solver = OptaxSolver(fun=to_minimize, opt=optax.adam(3e-4), implicit_diff=True)
initial, _ = solver.init(init_params = 5.)
result, _ = solver.run(init_params = initial)
return result
jax.value_and_grad(pipeline)(2., data=6.)
```
which yields this error:
```
CustomVJPException: Detected differentiation of a custom_vjp function with respect to a closed-over value. That isn't supported because the custom VJP rule only specifies how to differentiate the custom_vjp function with respect to explicit input parameters. Try passing the closed-over value into the custom_vjp function as an argument, and adapting the custom_vjp fwd and bwd rules.
```
My versions are:
```
jax==0.2.20
jaxlib==0.1.71
jaxopt==0.0.1
optax==0.0.9
```
Am I doing something very silly? I guess I'm also wondering if this example within the scope of the solver API? I noticed that this doesn't occur with `solver.update`, just with `solver.run`.
Thanks :)
Contributor guide
Assessment
This issue has not been assessed yet.