[linen] Linesearch (and lbfgs) support for TrainState

[emiresenov]

What does this PR do?

This PR modifies .update() to support the optional additional arguments value and value_fn for the GradientTransformationExtraArgs class.

These changes are in the same vein as the merged PR #4351, addressing #4144 for the TrainState class.

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).
• This change is discussed in a Github issue/
  Linesearch (and lbfgs) support #4351 , Support for optax lbfgs and related optimizers with NNX #4144
• The documentation and docstrings adhere to the
  documentation guidelines.
• This change includes necessary high-coverage tests.
  Linesearch (and lbfgs) support #4351 includes a test with these modifications. Additionally, the modifications follow the spec for GradientTransformationExtraArgs. If you want me to write a file with test cases for the TrainState class, let me know. Still, I cannot find any tests for this class in the repository.

[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]

I think if either of them is not None we should pass them.

[emiresenov]

Good idea. I'll update the PR

[emiresenov]

I figure the same logic should hold for grad since e.g. scale_by_polyak also usesGradientTransformationExtraArgs and only takes value as an additional keyword argument in .update().

I moved all three optional args to an update_kwargs dictionary, adding each to .update() only if they are not None. Additionally, my former commit should not have sent the **kwargs to .update() sinceTrainState uses these for something else – I mixed up the usage from the merged PR I was referring to.

See the changes in my latest commit and let me know if there are any issues.

[cgarciae]

This version looks better. Can you fix pre-commit issue? Run

pip install pre-commit
pre-commit run --all-files

[cgarciae]

Tangent: curious if nnx.Optimizer works with scale_by_polyak? I've never used it but we recently added some support for Linesearch in Optimizer as well.

[emiresenov]

This version looks better. Can you fix pre-commit issue? Run

pip install pre-commit
pre-commit run --all-files

Fixed, see my latest commit!

[emiresenov]

Tangent: curious if nnx.Optimizer works with scale_by_polyak? I've never used it but we recently added some support for Linesearch in Optimizer as well.

I haven't used it either, but I see no reason why it wouldn't work since the last commit of #4351 uses the same logic. I think the first commit of that PR (the one I mistakenly replicated in my first commit) suggested a solution that wouldn't work for scale_by_polyak used by optax.polyak_sgd since it didn't incorporate grad as an optional if only value was passed, presumably because the author only had Linesearch in mind.

However, this potential issue was fixed in the last commit of the PR when the author moved everything to their added kwargs in .update(). With this change, each additional value to the GradientTransformationExtraArgs update function is passed only if not None – this makes the solution more generalized and also enables scale_by_polyak which only takes value as an additional keyword argument.