value and grad from state but with aux outputs

[SNMS95]

I have an expensive loss function that returns a lot of intermediate results as auxiliary data.
I want to optimize it with LBFGS (with line-search). However, to cache the calculated gradients, I need to use
opt.value_and_grad_from_state but it does not allow to thread auxiliary values.

For a simple case, I have adapted the example into a MWE.

import optax
import jax.numpy as jnp
def fn(x): return jnp.sum(x ** 2), x

solver = optax.chain(
    optax.sgd(learning_rate=1.),
    optax.scale_by_backtracking_linesearch(
        max_backtracking_steps=15, store_grad=True
    )
)
value_and_grad = optax.value_and_grad_from_state(fn)
x = jnp.array([1., 2., 3.])
opt_state = solver.init(x)
value, grad = value_and_grad(x, state=opt_state)
print(value)
print(grad)

Is there an option like has_aux as is available in jax.value_and_grad

[rdyro]

I think it makes sense to add this option. Would you be interested in contributing a PR?

If you're interested in a quick workaround, optax has a very simple implementation of this functionality which you could directly reproduce in your code

optax/optax/_src/utils.py

Line 287 in 9b682ab

grad = otu.tree_get(state, 'grad')