Gradient through closure
- Dominant language
- Python
- Stars
- 1.1k
- Forks
- 76
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 1
Description
Is it possible to do the following?
I would like to take a gradient through a an argmin involving a closure, where the function passed to the solver contains `y`, which is not passed as an argument explicitly. This seems to be causing an issue.
```
from jaxopt import LBFGS
import jax.numpy as jnp
import jax
def implicit_layer(y):
def closure(x):
loss = jnp.sum(x**2) - jnp.sum(x * y)
return loss
lbfgs = LBFGS(fun=closure, tol=1e-5, stepsize=1e-1, maxiter=100, history_size=5,
use_gamma=True)
out, _ = lbfgs.run(y)
return out
jax.grad(implicit_layer)(y)
```
(this is just a toy example of what I have in mind)
The error message is:
```
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.
```
Contributor guide
Assessment
This issue has not been assessed yet.