JAX 0.6.0 Installation Instructions
Python 3.10.12
jaxlib
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp310-cp310-manylinux2014_x86_64.whl
JAX ROCm Plugin
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp310-cp310-manylinux_2_28_x86_64.whl
Python 3.11.0
jaxlib
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp311-cp311-manylinux2014_x86_64.whl
JAX ROCm Plugin
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp311-cp311-manylinux_2_28_x86_64.whl
Python 3.12.3
jaxlib
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp312-cp312-manylinux2014_x86_64.whl
JAX ROCm Plugin
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp312-cp312-manylinux_2_28_x86_64.whl
Python 3.13.8
jaxlib
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp313-cp313-manylinux2014_x86_64.whl
JAX ROCm Plugin
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp313-cp313-manylinux_2_28_x86_64.whl
jaxlib 0.6.2 Installation Instructions
Python 3.10.12
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp310-cp310-manylinux2014_x86_64.whl
Python 3.11.0
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp311-cp311-manylinux2014_x86_64.whl
Python 3.12.3
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp312-cp312-manylinux2014_x86_64.whl
Python 3.13.8
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp313-cp313-manylinux2014_x86_64.whl
What's Changed
- Add workflow for nightly pull from upstream by @charleshofer in #115
- Allow devs to kick off sync job manually by @charleshofer in #119
- Unpin container in CI build and remove libssl-dev install by @charleshofer in #123
- Rename the CI flow to 'ROCm CI' and only run it on PRs to rocm-main by @charleshofer in #126
- Fix nightly sync permissions by @charleshofer in #124
- Add GHA workflow for opening PRs upstream by @charleshofer in #116
- Fix upstream sync checkout by @charleshofer in #130
- Make daily sync permissions at the workflow level and fix merge CI by @charleshofer in #143
- Fix upstream PR workflow to use origin branches by @charleshofer in #151
- Add token to open upstream PRs by @charleshofer in #152
- Change the workflow for opening upstream PRs to post links that open PRs by @charleshofer in #157
- Fix rebase command to exclude rocm-main by @charleshofer in #158
- Fix user identity for rebase by @charleshofer in #159
- Fix the link to the downstream PR by @charleshofer in #160
- Use the reference format for links instead of inline by @charleshofer in #162
- Add upload wheels file for pypi by @Ruturaj4 in #184
- Run less tests for CI by @charleshofer in #199
- Run GPU CI on QA PRs by @charleshofer in #228
- Use bazel for PR tests by @charleshofer in #216
- removing pytest-csv from pytests by @zahiqbal in #248
- Add gfx1101 Target to Jax builds by @JehandadKhan in #249
- Use XLA with include fix by @charleshofer in #256
- Deprecate obsolete gfx versions by @JehandadKhan in #273
- Add JSON output to multi-GPU tests by @charleshofer in #274
- Count test totals correctly for dashboards by @charleshofer in #280
- Fix auditwheel version issue by @mrodden in #288
- Fix ROCm build README by @charleshofer in #284
- Add back raw totals in JSON reports by @charleshofer in #281
- Fixes for 0.5.0 build ported to rocm-main by @mrodden in #315
- Use rebase for upstream sync by @charleshofer in #325
- Add Python and ROCm version matrix for rocm-main by @charleshofer in #314
- adding jax unit test tridiagonal for rocm platform by @zahiqbal in #369
- [0.6.0-UT] Cherry pick UTs from 0.5.0 by @gulsumgudukbay in #438
- [0.6.0-UT] Add shared utility function get_rocm_version to test_util.py by @psanal35 in #445
- [0.6.0-UT] adding a QA test requirements file for the run_single_gpu test script… by @gulsumgudukbay in #446
- [0.6.0-UT] fixed typo by @gulsumgudukbay in #450
- Revert "[0.6.0-UT] fixed typo" and fix rocm-test-requirements by @gulsumgudukbay in #452
- [0.6.0-UT] Select the platform GPUs for a few MultiProcessGpuTest test cases. by @phambinhfin in #443
- [0.6.0-UT] Fix NumpyLinalgTest test cases by @phambinhfin in #453
- Build with ROCm 7.0 on new runner by @charleshofer in #449
- Fix get_rocm_version to skip tests if installation is missing by @psanal35 in #456
- Add missing header
warningsfor test_util.py by @psanal35 in #458 - Remove warnings.warn from get_rocm_version by @psanal35 in #459
- [0.6.0-UT] skip failing tests by @gulsumgudukbay in #442
- SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError by @gulsumgudukbay in #465
- [0.6.0-UT] Fix namespace rocm7 v2 by @gulsumgudukbay in #473
- [0.6.0-UT] Adjusting block sizes in BWD-FusedAttention tests to fit in GPU share… by @gulsumgudukbay in #490
- [0.6.0-UT] Fix test_gpu_distributed_initialize() to take number of available GPUs into account by @Arech8 in #501
- Fix test_gpu_distributed_initialize() to get a more proper GPUs number by @Arech8 in #503
- [0.6.0-UT] Fix ann_test::test_pmap by @gulsumgudukbay in #509
- Bugfix bazel debug rule by @Arech8 in #511
- Skip test_dot_precision() with TF32 for Navi architecture by @Arech8 in #520
- Fix GPU handle pool singleton aliasing by @phambinhfin in #525
- [0.6.0-UT] Fix is_device_rocm function calls by @gulsumgudukbay in #523
- [0.6.0-UT] Detect aborted tests and add abort info to UT logs by @gulsumgudukbay in #526
- [0.6.0-UT] Transferring Multi-GPU tests by @gulsumgudukbay in #527
- [0.6.0-UT] Adding abort support to run_multi_gpu by @gulsumgudukbay in #529
- Fix GPU handle pool singleton aliasing for HipSPARSE by @AratiGanesh in #535
- [0.6.0-UT] Unskipping passing tests by @gulsumgudukbay in #540
- [0.6.0-UT] Adding python equivalent of run_multi_gpu.sh by @gulsumgudukbay in #543
- [0.6.0-UT] Adding third field to return value of get_rocm_version by @gulsumgudukbay in #542
- Wrap HIP inline functions in anonymous namespaces in vendor.h by @AratiGanesh in #545
- Updating jaxlib BUILD file for new target and XLA commit hash by @gulsumgudukbay in #546
- [Fix] breaking BUILD file with rocm_hip rename by @GeneDer in #548
New Contributors
- @Arech8 made their first contribution in #501
- @AratiGanesh made their first contribution in #535
- @GeneDer made their first contribution in #548
Full Changelog: rocm-jax-v0.5.0...rocm-jax-v0.6.0
Contributors
mrodden, JehandadKhan, and 9 other contributors