google / google/jaxopt

Problem differentiating through `solver.run` in `OptaxSolver`

Open
#31 21 comments 3 reactions 0 assignees View on GitHub
question
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.