[hcb] Simplifications to the host_calback API #8678
Merged
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
hcb_simplify
branch
from
26707ff
to
9947929
Compare
gnecula
force-pushed
the
hcb_simplify
branch
3 times, most recently
from
d2ab70d
to
bd8236c
Compare
|
I guess the idea is that users should host_callback inside a custom_vjp if they want to capture values from the reverse pass? |
Indeed, that is the idea. In fact, some folks already do this because they want to pre-process the gradients on the device before sending them to the host. |
gnecula
force-pushed
the
hcb_simplify
branch
4 times, most recently
from
6ed9f07
to
a3f7040
Compare
gnecula
force-pushed
the
hcb_simplify
branch
from
a3f7040
to
deb292c
Compare
* dropping support for special AD handling for hcb.id_tap and id_print. From now on, only the primals are tapped. The old behavior can be obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS environment variale, or the --flax_host_callback_ad_transforms flag. Additionally, added documentation for how to implement the old behavior using JAX custom AD APIs. This allows us to make some significant cleanup in the internals.
gnecula
force-pushed
the
hcb_simplify
branch
from
deb292c
to
f08156a
Compare
copybara-service bot
pushed a commit
that referenced
this pull request
Aug 16, 2023
…ack_ad_transforms This flag was added in #8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. PiperOrigin-RevId: 557402331
copybara-service bot
pushed a commit
that referenced
this pull request
Aug 16, 2023
…back_ad_transforms. This flag was added in #8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. This is part of a longer project to replace uses of host_callback with jax.pure_callback and jax.experimental.io_callback. PiperOrigin-RevId: 557402331
copybara-service bot
pushed a commit
that referenced
this pull request
Aug 16, 2023
…back_ad_transforms. This flag was added in #8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. This is part of a longer project to replace uses of host_callback with jax.pure_callback and jax.experimental.io_callback. PiperOrigin-RevId: 557402331
copybara-service bot
pushed a commit
that referenced
this pull request
Aug 16, 2023
…back_ad_transforms. This flag was added in #8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. This is part of a longer project to replace uses of host_callback with jax.pure_callback and jax.experimental.io_callback. PiperOrigin-RevId: 557402331
copybara-service bot
pushed a commit
that referenced
this pull request
Aug 16, 2023
…back_ad_transforms. This flag was added in #8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. This is part of a longer project to replace uses of host_callback with jax.pure_callback and jax.experimental.io_callback. PiperOrigin-RevId: 557520668
NeilGirdhar
pushed a commit
to NeilGirdhar/jax
that referenced
this pull request
Aug 16, 2023
…back_ad_transforms. This flag was added in google#8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. This is part of a longer project to replace uses of host_callback with jax.pure_callback and jax.experimental.io_callback. PiperOrigin-RevId: 557520668
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.
From now on, only the primals are tapped.
This allows us to make some significant cleanup in the internals.