[better_errors] Add debug info to the Jaxprs formed for AD (step 2) #26313
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
force-pushed
the
debug_info_vjp
branch
2 times, most recently
from
32151e2 to
7606afc
Compare
gnecula
changed the title
gnecula
force-pushed
the
debug_info_vjp
branch
9 times, most recently
from
2331c9e to
5b9b852
Compare
dfm
approved these changes
Feb 5, 2025
| debug_fwd = debug_info("custom_vjp fwd", self.fwd, args, kwargs, | ||
| static_argnums=self.nondiff_argnums) | ||
| # TODO(necula): figure out how to construct the debug_bwd args | ||
| debug_bwd = debug_info("custom_vjp bwd", self.bwd, args, {}) |
There was a problem hiding this comment.
As noted in the TODO, this definitely isn't the right signature here. Would it be better to use a more obvious placeholder?
(Side note, while I'm here: I think the arguments to bwd are *static_args, residuals, output_cotangents, where residuals are output by fwd and output_cotangents has the tangent structure associated with fun's output. I can see why this will be tricky to annotate!)
There was a problem hiding this comment.
Indeed, this is tricky because I do not know the structure of the residuals and output_cotangents at this point.
The signatures are only needed to pick up the proper names for arg_names. Fixing that will have to come later, right now I am focusing just populating at least the function src info. We may never be able to get the arg_names perfectly right, there are so many places where functions are transformed. In that case, I will think (later) about a way to safely default to None. For now, the default is that we use None if the length of arg_names is different than jaxpr.invars.
Following jax-ml#26078 , we add debug info to more calls of lu.wrap_init.
gnecula
force-pushed
the
debug_info_vjp
branch
from
5b9b852 to
abcaec7
Compare
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 6, 2025
This follows after jax-ml#26078, and jax-ml#26313, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitives `custom_vjp_call_jaxpr` and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`, but in almost all cases if was really ` lu.WrappedFun.call_wrapped`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 7, 2025
This follows after jax-ml#26078, jax-ml#26313, jax-ml#26348, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitive `custom_transpose` to take the `transpose` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`. These changes ensure that all the `lu.wrap_init` and `Jaxpr` are called with debug_info in the `api_test.py:CustomTransposeTest`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 7, 2025
This follows after jax-ml#26078, jax-ml#26313, jax-ml#26348, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitive `custom_transpose` to take the `transpose` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`. These changes ensure that all the `lu.wrap_init` and `Jaxpr` are called with debug_info in the `api_test.py:CustomTransposeTest`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 7, 2025
This follows after jax-ml#26078, jax-ml#26313, jax-ml#26348, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitive `custom_transpose` to take the `transpose` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`. These changes ensure that all the `lu.wrap_init` and `Jaxpr` are called with debug_info in the `api_test.py:CustomTransposeTest`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 7, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. These changes ensure that all the lu.wrap_init and Jaxpr are called with debug_info in the api_test.py:CustomTransposeTest, api_test.py:CustomVmapTest and api_test.py:RematTest.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 7, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. These changes ensure that all the lu.wrap_init and Jaxpr are called with debug_info in the api_test.py:CustomTransposeTest, api_test.py:CustomVmapTest and api_test.py:RematTest.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 7, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. These changes ensure that all the lu.wrap_init and Jaxpr are called with debug_info in the api_test.py:CustomTransposeTest, api_test.py:CustomVmapTest and api_test.py:RematTest.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 7, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info).
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info).
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows after jax-ml#26078, jax-ml#26313, jax-ml#26348, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitive `custom_transpose` to take the `transpose` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`. These changes ensure that all the `lu.wrap_init` and `Jaxpr` are called with debug_info in the `api_test.py:CustomTransposeTest`.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. These changes ensure that all the lu.wrap_init and Jaxpr are called with debug_info in the api_test.py:CustomTransposeTest, api_test.py:CustomVmapTest and api_test.py:RematTest.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info).
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info).
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info). Also fixed uses in shard_map, checkify, and jax2tf.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info). Also fixed uses in shard_map, checkify, and jax2tf.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info). Also fixed uses in shard_map, checkify, and jax2tf.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info). Also fixed uses in shard_map, checkify, and jax2tf.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Fixes in jet, ...
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Fixes in jet, ...
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info). Also fixed uses in shard_map, checkify, sparse, attrs, and jax2tf.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Here I changed the `custom_jvp_call` to replace the parameter `jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun` that can carry debug info). Also fixed uses in shard_map, checkify, sparse, attrs, and jax2tf.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Fixes in jet, stateful code, tests.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Fixes in jet, stateful code, key_reuse, ode, tests.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 8, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Fixes in jet, stateful code, key_reuse, ode, tests.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 9, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Fixes in jet, stateful code, key_reuse, ode, pallas, tests.
gnecula
added a commit
to gnecula/jax
that referenced
this pull request
Feb 9, 2025
This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init. Fixes in jet, stateful code, key_reuse, ode, pallas, tests.
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.
Following #26078 , we add debug info to more calls of
lu.wrap_init.Most of the work here is to construct a debug info in appropriate places and thread it to places where
lu.wrap_initis called. There are just a lot of calls tolu.wrap_init, and this PR does not even fix all of them.We also add type declarations in all the functions that I had to understand for this PR.