google / google/jaxopt

Gradient through closure

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

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.