Linesearch (and lbfgs) support

[jlperla]

What does this PR do?

This PR modifies the .update() to support additional arguments required for the Optax lbfgs and related algorithms which rely on linesearch methods.

Fixes #4144

• Additional tests were added to check the lbfgs() (where I believe the additional arguments are required for the linesearch more generally)
• The key feature needed is a callback to evaluate the objective function at different "state" values. This needs to use a split and merge in the optimization.
  • In cases where the model_static, model_state = nnx.split(state.model, self.wrt) has already been done, the model_static can be passed into the .update to avoid doing the nnx.split step with each iteration. See the test_jit_linesearch for more there.

A feature missing here is support for value_and_grad_from_state as described in https://optax.readthedocs.io/en/stable/_collections/examples/lbfgs.html#linesearches-in-practice. Implementing this would provide some performance advantages as it would allow the optimizer to reuse the gradients/value.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[cgarciae]

Hey @jlperla! I left a longer comment in #4144 but TLDR is that I think we should either only add **kwargs and have the user implement the definition for value_fn, or implement an entirely new optimizer class for this family of algorithms (or just specific to lbfgs).

[jlperla]

@cgarciae See if this is what you had in mind. If so, it seems like it solves the general GradientTransformationExtraArgs challenges for future features. There are other arguments there as well.

[cgarciae]

@jlperla sound reasonable. Approved.

[jlperla]

Amazing, thanks @cgarciae When is the next expected release? Would love to publicize this in some sample code for a paper.