Fix parameter-shift bug in JAX interface #2255
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
Codecov Report
@@ Coverage Diff @@
## master #2255 +/- ##
=======================================
Coverage 99.27% 99.27%
=======================================
Files 232 232
Lines 18712 18713 +1
=======================================
+ Hits 18577 18578 +1
Misses 135 135
Continue to review full report at Codecov.
|
|
@antalszava I can't think of a good test to add to ensure this is fixed 🤔 |
antalszava
reviewed
Mar 2, 2022
| @@ -427,6 +429,7 @@ def cost_fn(a, b): | |||
|
|
|||
| expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] | |||
| assert np.allclose(res, expected, atol=0.1, rtol=0) | |||
| assert all(not isinstance(p, jnp.ndarray) for p in spy.call_args[0][0][0].get_parameters()) | |||
There was a problem hiding this comment.
@josh146 this would fail on current master 👍
There was a problem hiding this comment.
very nice!! Thanks @antalszava, this is a smart way of testing it
|
[sc-15391] |
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.
Context: Applying the parameter-shift rule with the JAX interface was failing, as VJP tapes sent for execution were not unwrapping their parameters.
Description of the Change: Fixes the bug by unwrapping VJP tapes prior to execution.
Benefits: Works as expected on devices that do not support
DeviceArrayobjects.Possible Drawbacks: This will be a blocker to supporting higher-order derivatives with JAX, since we are converting device array's to NumPy arrays. Although I don't believe this was supported previously yet anyway.
Related GitHub Issues: Fixes #2242