Consolidate initial/final style custom_vjp primitives into one #28589
Conversation
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
dfm
force-pushed
the
consolidate-custom-vjp-primitives
branch
2 times, most recently
from
9ed9482 to
8941ee8
Compare
dfm
force-pushed
the
consolidate-custom-vjp-primitives
branch
4 times, most recently
from
df642f1 to
a0cf239
Compare
dfm
force-pushed
the
consolidate-custom-vjp-primitives
branch
4 times, most recently
from
0132a9c to
883e49b
Compare
copybara-service bot
pushed a commit
to google/aqt
that referenced
this pull request
May 12, 2025
In jax-ml/jax#28589, we're slightly tweaking the behavior of custom_vjp when staged out. Of relevance here, the Jaxpr parameter is now called `call_jaxpr` (for consistency with other internal higher-order primitives) instead of `fun_jaxpr`. This change supports old and new versions of JAX. PiperOrigin-RevId: 757772896
copybara-service bot
pushed a commit
to google/aqt
that referenced
this pull request
May 12, 2025
In jax-ml/jax#28589, we're slightly tweaking the behavior of custom_vjp when staged out. Of relevance here, the Jaxpr parameter is now called `call_jaxpr` (for consistency with other internal higher-order primitives) instead of `fun_jaxpr`. This change supports old and new versions of JAX. PiperOrigin-RevId: 757772896
copybara-service bot
pushed a commit
to google/aqt
that referenced
this pull request
May 12, 2025
In jax-ml/jax#28589, we're slightly tweaking the behavior of custom_vjp when staged out. Of relevance here, the Jaxpr parameter is now called `call_jaxpr` (for consistency with other internal higher-order primitives) instead of `fun_jaxpr`. This change supports old and new versions of JAX. PiperOrigin-RevId: 757883900
dfm
force-pushed
the
consolidate-custom-vjp-primitives
branch
from
883e49b to
74938be
Compare
copybara-service bot
pushed a commit
to jax-ml/jax-tpu-embedding
that referenced
this pull request
May 15, 2025
The primitive name is changed in jax-ml/jax#28589 PiperOrigin-RevId: 759252787
copybara-service bot
pushed a commit
to jax-ml/jax-tpu-embedding
that referenced
this pull request
May 15, 2025
The primitive name is changed in jax-ml/jax#28589 PiperOrigin-RevId: 759268475
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.
Having separate primitives for initial and final style meant that we needed some duplicated logic so this should be a net win.
Most of the tests pass, but there are a couple of failures related to closed-over tracers.All tests green after #28605 goes in.