Lower jax.numpy matmul functions to mixed-precision dot_general #16736
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
froystig
approved these changes
Jul 17, 2023
jakevdp
force-pushed
the
jnp-mixed-precision
branch
2 times, most recently
from
cbec998
to
ac6410d
Compare
jakevdp
force-pushed
the
jnp-mixed-precision
branch
3 times, most recently
from
c7517f4
to
ade4b5a
Compare
jakevdp
force-pushed
the
jnp-mixed-precision
branch
from
ade4b5a
to
d8478c1
Compare
jakevdp
added
pull ready
jakevdp
force-pushed
the
jnp-mixed-precision
branch
3 times, most recently
from
66f03bc
to
dd71006
Compare
jakevdp
force-pushed
the
jnp-mixed-precision
branch
from
dd71006
to
50f3f2b
Compare
jakevdp
force-pushed
the
jnp-mixed-precision
branch
3 times, most recently
from
ad93bd4
to
f999827
Compare
jakevdp
force-pushed
the
jnp-mixed-precision
branch
from
f999827
to
7d29ed6
Compare
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.
Related to #16703
I opted not to add specific tests for these as I did for
dotin #16721 andeinsumin #16733Note that the sparse test tolerances needed to change because the new path no longer calls
convert_element_type, which for sparse inputs has a side-effect of callingsum_duplicates, which previously improved the numerical accuracy of the matmul output at the expense of a slower runtime. I opted not to force asum_duplicatesinbcoo_dot_general, but rather to allow for this inaccuracy in order to improve performance of the default impl.