[better_errors] Ensure that tracer errors in for_loop point to use code #25798
+36
−6
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
gnecula
changed the title
gnecula
force-pushed
the
fix_fori_error
branch
from
844b74a to
231ddbf
Compare
jakevdp
reviewed
Jan 10, 2025
gnecula
force-pushed
the
fix_fori_error
branch
2 times, most recently
from
1f0ed9f to
7bdeb19
Compare
jakevdp
reviewed
Jan 13, 2025
jax/_src/lax/control_flow/loops.py
Outdated
| body_fun = weakref.ref(body_fun) | ||
| body_fun_ref = weakref.ref(body_fun) | ||
|
|
||
| @wraps(body_fun) |
There was a problem hiding this comment.
Actually, looking more closely here, we probably can't use this mechanism (either functools.wraps or functools.update_wrapper) because it results in a strong reference of the original function being attached to the wrapped function: https://github.com/python/cpython/blob/aa6579cb60b4fcd652d7bc83fed2668e4ae84db3/Lib/functools.py#L61
There was a problem hiding this comment.
You are right, then the solution is to inline (parts of) wraps to get the function name to work properly.
There was a problem hiding this comment.
In the end, I decided to just save the fun_sourceinfo, rather than saving all the ingredients (the filename and lineno are read-only).
gnecula
force-pushed
the
fix_fori_error
branch
from
7bdeb19 to
3faff78
Compare
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.
Fixes: #23637