jax

c2a29f63 - Merge pull request #34600 from ROCm:phambinh/ci_rocm_scaled_matmul_tests