jax

8e541a61 - adding run_pytest_rocm.sh