Update optimization to use rrule system #23
Merged
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
lkdvos
commented
May 1, 2024
pbrehmer
previously approved these changes
May 2, 2024
There was a problem hiding this comment.
The refactor looks really great, overall it's a big improvement. Thanks!
The gradient test was broken since it wasn't yet adapted to the hook_pullback wrapper, so I quickly repaired it. Also, the ManualIter fixed-point differentiation had a weird convergence criterion. I changed that to make it comparable to GeometricSum (these two approaches are kind of equivalent anyways).
|
Great, I'll look into the documentation pr first, as this is probably easier to merge in that order? |
|
Yes that's a good idea, I would think it is easier to merge that way. |
|
I'll rebase tomorrow and get this ready to merge |
lkdvos
force-pushed
the
lb/leading_boundary_diff
branch
from
9086899
to
bcb29fc
Compare
pbrehmer
approved these changes
May 3, 2024
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.
This is a major refactor of the optimization code to have a uniform interface.
In particular, I implemented all of the
ctmrg_gradientfunctions through AD by using @pbrehmer 's implementation of the gradient and transforming them into rrules via @leburgel 's adaptation.Then, in order to be able to dynamically select which rrule to use, I added the function
hook_pullback, which functions as a convenient interface to change the rrule algorithm through an optional keyword that can later be dispatched on in the rrule.Finally, I added some rrules for constructors/getters for InfinitePEPS and CTMRGEnv, which prevents our costfunction from yielding weird tangent types. This liquidates lots of the checks to see if the tangents are the correct type and gives a bit of a cleaner implementation. I have not played around with how robust this is, and if custom costfunctions would easily break this, but I did add a type check in the implementation of the rrules to at least detect when this is happening.