jax

fe4bbbc0 - Merge pull request #34500 from magaonka-amd:fix-keyerror-rocm-upstream