[QwixMxfpNumerics] Verify XLA handling of MXFP data types by testing jax.nn.scaled_matmul across different platforms.#227
Merged
copybara-service[bot] merged 1 commit intomainfrom Mar 23, 2026
Conversation
copybara-service
Bot
force-pushed
the
test_881171796
branch
from
2fdf9f2 to
eaf7eae
Compare
copybara-service
Bot
force-pushed
the
test_881171796
branch
2 times, most recently
from
08fccde to
0da0db4
Compare
…`jax.nn.scaled_matmul` across different platforms. This change introduces `mxfp_numerics_test.py` to validate the behavior of `jax.nn.scaled_matmul` with mixed-precision floating-point types (MXFP4 and MXFP8). The tests cover a float32 baseline, MXFP8 using e4m3fn/e8m0fnu, and MXFP4 using e2m1fn/e8m0fnu. The results are compared against a float32 reference implementation. HLO is logged to check for native hardware support. BUILD rules are added to run these tests on CPU, TPUs(v5e and v6e), and GPUs(H100 and B200). PiperOrigin-RevId: 888251069
copybara-service
Bot
force-pushed
the
test_881171796
branch
from
0da0db4 to
e43873c
Compare
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
[QwixMxfpNumerics] Verify XLA handling of MXFP data types by testing
jax.nn.scaled_matmulacross different platforms.This change introduces
mxfp_numerics_test.pyto validate the behavior ofjax.nn.scaled_matmulwith mixed-precision floating-point types (MXFP4 and MXFP8). The tests cover a float32 baseline, MXFP8 using e4m3fn/e8m0fnu, and MXFP4 using e2m1fn/e8m0fnu. The results are compared against a float32 reference implementation. HLO is logged to check for native hardware support. BUILD rules are added to run these tests on CPU, TPUs(v5e and v6e), and GPUs(H100 and B200).