fix custom_jvp/vjp closure issues, and nondiff_argnums too! #4008
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
mattjj
force-pushed
the
custom-jvp-closure-fixes
branch
2 times, most recently
from
b215d2f
to
070beae
Compare
mattjj
changed the title
This was referenced Oct 12, 2020
mattjj
force-pushed
the
custom-jvp-closure-fixes
branch
from
4299e30
to
8f51e6a
Compare
copybara-service bot
pushed a commit
to google-deepmind/dm-haiku
that referenced
this pull request
Oct 14, 2020
Currently this method is inherited from the Trace base class in jax's core.py, but I want to remove that default method because falling through to it has caused at least two bugs (one in jax2tf that came up in jax-ml/jax#4008, and another at HEAD described here: jax-ml/jax#4566). This won't change Haiku's behavior at all, but it'll make sure things keep working when the methods on the base Trace class disappear. PiperOrigin-RevId: 337088168 Change-Id: I4d5f85e6a2db5fc6348a2f17445afd23d29d9c9c
mattjj
added a commit
that referenced
this pull request
Oct 15, 2020
This change sets up some internal users so that we can then land #4008.
mattjj
added a commit
that referenced
this pull request
Oct 15, 2020
This caused a test failure when trying to land #4008.
copybara-service bot
pushed a commit
to tensorflow/probability
that referenced
this pull request
Oct 16, 2020
HarvestTrace.process_custom_vjp_call. Currently these fall back to the defaults on the superclass Trace, defined in jax/core.py, but those defaults are going away after jax-ml/jax#4008. Moreover, they probably didn't have the behavior we want! PiperOrigin-RevId: 337426371
copybara-service bot
pushed a commit
to google/trax
that referenced
this pull request
Oct 16, 2020
jax-ml/jax#4008. In particular, don't use nondiff_argnums for array-valued arguments to custom_vjp functions. This change broke three test cases, which are now skipped. I'll work with lukaszkaiser@ to fix them! PiperOrigin-RevId: 337338931
mattjj
added a commit
that referenced
this pull request
Oct 16, 2020
copybara-service bot
pushed a commit
to google/trax
that referenced
this pull request
Oct 16, 2020
jax-ml/jax#4008. In particular, don't use nondiff_argnums for array-valued arguments to custom_vjp functions. This change broke three test cases, which are now skipped. I'll work with lukaszkaiser@ to fix them! PiperOrigin-RevId: 337338931
copybara-service bot
pushed a commit
to google/trax
that referenced
this pull request
Oct 16, 2020
jax-ml/jax#4008. In particular, don't use nondiff_argnums for array-valued arguments to custom_vjp functions. This change broke three test cases, which are now skipped. I'll work with lukaszkaiser@ to fix them! PiperOrigin-RevId: 337451477
mattjj
force-pushed
the
custom-jvp-closure-fixes
branch
from
af9d7a1
to
609f6f3
Compare
copybara-service bot
pushed a commit
to tensorflow/probability
that referenced
this pull request
Oct 16, 2020
PiperOrigin-RevId: 337538591
apaszke
reviewed
Oct 19, 2020
|
|
||
| def clip_gradient_bwd(res, g): | ||
| lo, hi = res | ||
| return (None, None, jnp.clip(g, lo, hi)) # return None for lo and hi |
There was a problem hiding this comment.
There's interesting point here: when an argument can sometimes be a tracer, but sometimes will just be a constant, and you're not actually tracing your code, you might want to query whether returning None is ok. The alternative is to always put in the work to compute the value, even though it might get thrown away later. But idk how important eager mode is for JAX, so it might not matter too much (it definitely did in PyTorch) 🤷
There was a problem hiding this comment.
Hmm, I may not be following: returning None is always okay here, as it just is shorthand for "zeros like the input". That is, it works regardless of whether the corresponding input was a tracer, and works efficiently both in eager mode and in jit mode. Can you say more?
|
For posterity: this was merged, as 4a20eea, not sure why copybara got confused. |
|
Woohoo, this is huge! |
This was referenced Oct 22, 2020
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 #4566, fixes #4521, fixes #2912, fixes #3822, fixes #4173, fixes #2520, fixes #3808
If you survey those issues, you'll come to the conclusion that
jax.custom_jvpandjax.custom_vjpdon't work with closing over JAX Tracers, and moreover thatnondiff_argnumscan be just as problematic. Indeed,nondiff_argnumswas effectively implemented in terms of lexical closure. The original plan was not to makecustom_jvp/custom_vjpwork with lexical closure, but I never documented that, and the error messages were terrible!This PR makes
custom_jvpandcustom_vjpwork with closed-over Tracers. Woo! That is, nowcustom_jvpandcustom_vjpfunctions and rules can close over Tracers to our hearts' content. For all non-autodiff transformations, things will Just Work. For autodiff transformations, we'll get a clear error message about why we can't differentiate with respect to values over which a custom_jvp or custom_vjp closes:This PR accomplishes that goal by following through on the original ansatz of
custom_jvpandcustom_vjp, namely to make them work likecore.call. That is, we do all thecore.process_env_tracesstuff properly.TODO: