jax
cdb399a1 - Remove version checks for ml_dtypes >= 0.5.0

Commit

188 days ago

Remove version checks for ml_dtypes >= 0.5.0 This PR removes runtime checks for ml_dtypes >= 0.5.0 related to float8_e8m0fnu and float4_e2m1fn, since version 0.5.0 is now the minimum required version for JAX. The ml_dtype requirement in JAX was updated to version 0.5.0+ (on Mar 20, 2025) - commit 4b7ead4 - The checks are no longer necessary and simplify the test logic. - All supported environments are guaranteed to have ml_dtypes >= 0.5.0. PiperOrigin-RevId: 783511371

References

#30235 - Remove version checks for ml_dtypes >= 0.5.0

#31381 - Remove old ROCm build code

#31720 - Fix ann_test.py numerical bug in target reshape

#31768 - [ROCm] Support lowering through PJRT_Triton_Extension

#32115 - Relax version requirements for ROCm Jax Plugin wheels

#33157 - Resolve undefined behavior in bitshift unit test

#579 - Create rocm-test-requirements.txt

#581 - Fix/pallas tests shared memory

#584 - Use plain bazel to test jax, use hermetic rocm dependency

#585 - update a test for checking zero ROCm GPU event

#34135 - [ROCm] update to test if there are GPU events when doing profiling on…