github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-02-13 07:55:07 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add JAX Plugin-PJRT support table

Browse Source
This commit is contained in:
Istvan Kiss
2025-11-17 11:37:37 +01:00
parent 52f6f0cbc4
commit bb84806ffe
1 changed files with 17 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
docs/compatibility/ml-compatibility/jax-compatibility.rst
View File

@@ -47,6 +47,23 @@ with ROCm support:
`Community ROCm JAX Docker images <https://hub.docker.com/r/rocm/jax-community>`_
follow upstream JAX releases and use the latest available ROCm version.
JAX Plugin-PJRT with JAX/JAXLIB compatibility
================================================================================
Portable JIT Runtime (PJRT) is an open, stable interface for device runtime and
compiler. The following table details the ROCm version compatibility matrix
between JAX Plugin–PJRT and JAX/JAXLIB.
.. list-table::
:header-rows: 1
* - JAX Plugin-PJRT
- JAX/JAXLIB
- ROCm
* - 0.6.0
- 0.6.2, 0.6.0
- 7.0.0, 7.0.1, 7.0.2
Use cases and recommendations
================================================================================