jax

08d80110 - Merge pull request #34602 from ROCm:phambinh/ci_rocm_pallas_ops_tests