Fixing gather batching rule not implemented when multiple qjits call them
#1096
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
…lling it out into patches.py as a global. This way, jax will not complain about "gather batching rule not find", since we patch with the same gather2_p object every qjit run.
joeycarter
approved these changes
Sep 4, 2024
erick-xanadu
approved these changes
Sep 4, 2024
There was a problem hiding this comment.
You are just moving the standard_primitive to the global scope, right? That sounds good. Do you know if this also fixes these tests?
test/pytest/test_callback.py:@pytest.mark.parametrize("order", ["good", "bad"])
test/pytest/test_callback.py:def test_vjp_as_residual(arg, order):
test/pytest/test_callback.py: if order == "bad":
test/pytest/test_callback.py: if order == "bad":
test/pytest/test_callback.py:@pytest.mark.parametrize("order", ["good", "bad"])
test/pytest/test_callback.py:def test_vjp_as_residual_automatic(arg, order):
test/pytest/test_callback.py: if order == "bad":
test/pytest/test_callback.py: if order == "bad":
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #1096 +/- ##
=======================================
Coverage 97.65% 97.65%
=======================================
Files 76 76
Lines 10769 10769
Branches 1245 1245
=======================================
Hits 10517 10517
Misses 203 203
Partials 49 49 ☔ View full report in Codecov by Sentry. |
just checked, these are fixed as well |
|
@paul0403 if you want, you can remove the skips in this PR or in another one that follows. Happy with either. |
I'll just unskip them here. |
rauletorresc
pushed a commit
that referenced
this pull request
Sep 11, 2024
…l them (#1096) **Context:** There is a bug if multiple qjits call into `gather` primitive #1094 . With how we currently patch them, jax retains the first custom `gather2_p` object and keeps using that in the second run. **Description of the Change:** Making sure there is only ever one `gather2_p` primitive object by pulling it out into patches.py as a global. This way, jax will not complain about "gather batching rule not find", since we patch with the same gather2_p object every qjit run. **Benefits:** Multiple qjits in the same program can call into gather primitives. Although in the standard recommended usage there is just one qjit in the overall entry function, this unblocks potential pytests (where multiple qjits cannot be avoided). **Related GitHub Issues:** closes #1094 , closes #894 [sc-72775]
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:
There is a bug if multiple qjits call into
gatherprimitive #1094 . With how we currently patch them, jax retains the first customgather2_pobject and keeps using that in the second run.Description of the Change:
Making sure there is only ever one
gather2_pprimitive object by pulling it out into patches.py as a global.This way, jax will not complain about "gather batching rule not find", since we patch with the same gather2_p object every qjit run.
Benefits:
Multiple qjits in the same program can call into gather primitives. Although in the standard recommended usage there is just one qjit in the overall entry function, this unblocks potential pytests (where multiple qjits cannot be avoided).
Related GitHub Issues: closes #1094 , closes #894
[sc-72775]