Jaxopt integration

[ronceray]

Hello,
I am interested in using general nonlinear optimizers (CG, BFGS...) in a jax framework. Jaxopt would seem much more adapted than optax (simpler & not targeted towards neural nets) but their github page says they are merging with optax. Is this happening reasonably soon? Can we expect the features of Jaxopt to be all ported to the next version of optax?
[independent, side question: is there a way to access the Hessian of the objective function in optax? I haven't found one.]

[rdyro]

The migration is still in progress, for the derivatives, optax includes wrappers around JAX autodiff, so would jax.hessian work for you? If you compute the value and hessian in the same JIT function you make your own optimized version of jax.value_and_grad

[ronceray]

Thanks for your reply. So I take it that there is no planned date for the migration. I'll work with Jaxopt for now then.
Regarding the Hessian - yes, of course, thanks. But this is computationally costly in my case. I should have been more specific: I was wondering whether you could access the hessian approximation as a byproduct of optimizers (as is the case with scipy.minimize for instance through the 'hess': https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.OptimizeResult.html#scipy.optimize.OptimizeResult).
I guess essentially what I'm looking for is a jax implementation of scipy.minimize!

[rdyro]

Makes sense! For example, for the main quasi-second order optimizer in optax L-BFGS, the inverse Hessian is never really constructed explicitly, but you can reconstruct it from the optimizer state, for L-BFGS:

optax/optax/_src/transform.py

Lines 1448 to 1470 in bff9977

	class ScaleByLBFGSState(NamedTuple):
	"""State for LBFGS solver.
	Attributes:
	count: iteration of the algorithm.
	params: current parameters.
	updates: current updates.
	diff_params_memory: represents a list of past parameters' differences up to
	some predetermined ``memory_size`` fixed in :func:`optax.scale_by_lbfgs`.
	diff_updates_memory: represents a list of past gradients/updates'
	differences up to some predetermined ``memory_size`` fixed in
	:func:`optax.scale_by_lbfgs`.
	weights_memory: list of past weights multiplying the rank one matrices
	defining the inverse Hessian approximation, see
	:func:`optax.scale_by_lbfgs` for more details.
	"""
	count: chex.Numeric
	params: base.Params
	updates: base.Params
	diff_params_memory: chex.ArrayTree
	diff_updates_memory: chex.ArrayTree
	weights_memory: chex.Array

optax doesn't currently have the same functionality as scipy.minimize, but if you can provide scipy.minimize with fast jit-ed JAX functions, you can probably make it run really fast!

[ronceray]

I feel stupid for not having thought of this. Just providing jitted functions to scipy.minimize does the job perfectly for me. Thanks a lot for your answer and help!

[rdyro]

Not at all, I'm glad it worked for you!