jax

4c157782 - Merge pull request #34571 from magaonka-amd:enable-reduce-window-tests-rocm-upstream