Remove conservative checks for supported Subtensors operations in JAX #849
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
ricardoV94
commented
Jun 24, 2024
| out_pt = x_pt[1, 2, 0] | ||
| assert isinstance(out_pt.owner.op, pt_subtensor.Subtensor) | ||
| out_fg = FunctionGraph([], [out_pt]) | ||
| compare_jax_and_py(out_fg, []) | ||
| out_fg = FunctionGraph([x_pt], [out_pt]) |
There was a problem hiding this comment.
Otherwise I suspect we were not really testing jax subtensor dispatch as the index operation would be just constant_folded. This was not always the case, as we used to not run any rewrites in compare_jax_and_py. Now we do
ricardoV94
commented
Jun 24, 2024
| compare_jax_and_py(out_fg, [x_np]) | ||
|
|
||
| # Boolean indexing should work if indexes are constant | ||
| out_pt = x_pt[np.random.binomial(1, 0.5, size=(3, 4, 5))] |
There was a problem hiding this comment.
This was the case that led me to open this PR. JAX is happy to do it but our check wouldn't let it compile
ricardoV94
force-pushed
the
jax_eager_subtensor
branch
from
c27898a
to
6fc8f7a
Compare
The check was failing incorrectly for cases that are supported such as constant Boolean arrays. Besides that, user may dispatch without necessarily jitting the graph. There is no reason to fail eagerly.
ricardoV94
force-pushed
the
jax_eager_subtensor
branch
from
6fc8f7a
to
2b99224
Compare
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #849 +/- ##
==========================================
- Coverage 80.89% 80.88% -0.01%
==========================================
Files 169 169
Lines 46979 46966 -13
Branches 11478 11472 -6
==========================================
- Hits 38002 37989 -13
Misses 6764 6764
Partials 2213 2213
|
ricardoV94
changed the title
lucianopaz
approved these changes
Jun 28, 2024
11 tasks
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.
Description
The check was failing incorrectly for cases that are supported such as constant Boolean arrays. Besides that, user may dispatch without necessarily jitting the graph. There is no reason to fail eagerly.
Related Issue
Checklist
Type of change