Adds an option to support forward-mode automatic differentiation in all minimisers #114
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
added 2 commits
February 2, 2025 13:14
… all minimisers. Moves gradient computation, handling and related documentation into lin_to_grad. Also removes an old jacobian function that used size-based heuristics to toggle between modes.
|
This addresses #112 (comment). |
|
Summary:
|
|
Tests were breaking because pyright now complains about the return type of |
Comment on lines
+207
to
+208
| if isinstance(solver, optx.OptaxMinimiser): # No support for forward option | ||
| return |
There was a problem hiding this comment.
Nit: it'd be better to do
@pytest.mark.parametrize("solver", [solver for solver in minimisers if not isinstance(solver, optx.OptaxMinimiser)]
Not important enough to block merging, just noting in passing!
There was a problem hiding this comment.
Noted for the future :)
|
Awesome! Merged :D |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Minimisers now accept a
"fwd"and a"bwd"option, andlin_to_gradwill use transposition to get the gradient with reverse-mode, while calling a custom_jacfwdthat evaluates the linearised function with unit-pytrees when using forward-mode. Where relevant, gradient-computation related documentation has been moved to these two functions.I have parameterised the tests for
minimiseandminimise_jvpto try both options, let me know if there are other tests I should also parameterise. I'm assuming everything else will also work if these two pass.I have also removed a function
jacobianfrom the_miscmodule, that used size-based heuristics to toggle between forward- and reverse-mode automatic differentiation. Since we now support explicit options in all solvers that compute derivatives, this function no longer fits our current way of doing things. (I don't think it was used anywhere anyway.)