Cherry-pick fixes for JAX 0.9.2 release from JAX 0.9.1 #740
magaonka-amd
merged
5 commits into
ROCm:rocm-jaxlib-v0.9.2
from
magaonka-amd:rocm-jaxlib-v0.9.2
Skip ann_test test_pmap on ROCm due to IndivisibleError
5661c9f1
Skip test_tridiagonal_solve_grad0 on AMD due to rocSparse numerical i…
17b99081
Add ROCm support for scaled matrix multiplication and scaled dot lowe…
bb0c1d6a
Clarify ROCm scaled_matmul lowering rationale.
3999c9fa
push ci
e1b1128b
magaonka-amd
changed the title 
gulsumgudukbay
approved these changes
on 2026-03-23
Login to write a write a comment.
Login via GitHub