jax
fc74a0aa - [pallas] Relax check for block spec equality in fuser.

Commit

186 days ago

[pallas] Relax check for block spec equality in fuser. Most functions in Python use identity equality, so the block spec equality will fail even when the index map functions are equivalent (e.g. from closures or `partial`). PiperOrigin-RevId: 776397038

References

#29766 - [pallas] Relax check for block spec equality in fuser.

#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

#33186 - Make nvidia version data optional for ROCm builds

#579 - Create rocm-test-requirements.txt

#580 - GESVDJ support for ROCm GPUs in JAX

#581 - Fix/pallas tests shared memory

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