jax

7cb36e57 - Merge pull request #36851 from magaonka-amd:feature/rocm-tpu-style-pytest