How to access objective function value?

[pvasired]

Hi,

I have an optimization loop where I run something like:

grads = update()
updates, opt_state = optimizer.update(grads, opt_state, params=T)
T = optax.apply_updates(T, updates)

I am wondering if anywhere in this loop I have access to the objective function value. Right now I externally calculate this objective function value so that I can plot the loss curve after the optimization has terminated. The function calls are somewhat costly, so I am just wondering if optax already has access to these values at every step.

[ramithuh]

I think you might be needing jax.value_and_grad in your update() function