Fix Scan higher order derivatives#1975
Merged
Merged
Conversation
ricardoV94
force-pushed
the
scan_3rd_derivative_bug
branch
3 times, most recently
from
789ee32 to
1b15cf6
Compare
ricardoV94
force-pushed
the
scan_3rd_derivative_bug
branch
from
1b15cf6 to
55fab8a
Compare
jessegrabowski
approved these changes
Mar 26, 2026
| # The gradient of an overwrite must zero out the direct pass-through from the | ||
| # old value; the only gradient path is through the output expression that replaced | ||
| # it (already captured by compute_all_gradients via known_grads). | ||
| overlapping_taps = set() |
Member
There was a problem hiding this comment.
Is this section the actual fix?
ricardoV94
commented
Mar 26, 2026
Comment on lines
+2712
to
+2713
| if dx in overlapping_taps: | ||
| continue # gradient truncates here |
Member
Author
There was a problem hiding this comment.
This is the fix, skipping the accumulation
ricardoV94
force-pushed
the
scan_3rd_derivative_bug
branch
from
55fab8a to
1e86a3a
Compare
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Closes #1772
I let Claude take a stab at this as it was not priority (broken since ever?). It's a bit over my head because if there's anything more complex than the ScanOp, it's the gradient of the ScanOp (derivatives are simpler anyone?).
Here is the best AI guess as to how it is organized: https://gist.github.com/ricardoV94/809a2c25ec61d149209336ab0add6f30. I corrected it several times, so take it with a big grain of salt.
Anyway it sounds like the gradient of a mit-mot output was wrong. Ignoring, where they come from, a mit-mot represents a tape that is read and written into multiple positions at each step. The reverse pass of this process should not flow through the written positions into the past (the same way that the reverse pass of x[0].set(y), wrt to x truncates at 0).
And apparently (or so Claude tried to persuade me), this truncation was not happening. The error only shows up if you differentiate more than twice, don't ask me why.
I won't spend more time understanding the L_op right now, so I would rest on the fact that it makes some conceptual sense and that the tests now pass.
And some sparse documentation on L_op and related functionality