[better_errors] Add debug_info to DynamicJaxprTrace and JaxprStackFrame #25827
Conversation
gnecula
force-pushed
the
debug_info_2
branch
6 times, most recently
from
a457184 to
4baa6f2
Compare
| @@ -83,9 +83,10 @@ def jvpfun(f, instantiate, transform_stack, primals, tangents): | |||
| return out_primals, out_tangents | |||
|
|
|||
| @lu.transformation_with_aux2 | |||
| def linearize_subtrace(_f, _store, _tag, nzs_in, *primals, **params): | |||
| def linearize_subtrace(_f: Callable, _store, _tag, nzs_in, *primals, **params): | |||
There was a problem hiding this comment.
Shouldn't the type of _f be lu.WrappedFun?
There was a problem hiding this comment.
Actually, no, it should be a Callable. (I added it, because I had the same confusion earlier.) See how below we invoke it, and WrappedFun is not a Callable.
There was a problem hiding this comment.
I don't think that's quite right! linearize_subtrace is called in LinearizeTrace.process_call, which means that the _f argument is the fun from call_p, which (for better or worse) is a WrappedFun.
Also, as far as I know, the first input to a linear_util transformation must always be a WrappedFun!
There was a problem hiding this comment.
Edit: Ignore what I said above! I understand what you mean now. In the new version of linear_util without context the inputs are magically Callables (and I now parsed your explanation properly). Sorry for my confusion!
gnecula
force-pushed
the
debug_info_2
branch
2 times, most recently
from
c47b287 to
d0f6b66
Compare
This is part of a sequence of changes to ensure that the debugging information is propagated properly. Additional cleanup: * Rename `result_paths` to `result_paths_thunk` in `TracingDebugInfo` to clarify the difference from the similar field in `JaxprDebugInfo` * Added more type declarations
This is part of a sequence of changes to ensure that the debugging information is propagated properly.
Additional cleanup:
result_pathstoresult_paths_thunkinTracingDebugInfoto clarify the difference from the similar field inJaxprDebugInfo