diff --git a/docs/compatibility/ml-compatibility/jax-compatibility.rst b/docs/compatibility/ml-compatibility/jax-compatibility.rst index 479256c48..33ab23685 100644 --- a/docs/compatibility/ml-compatibility/jax-compatibility.rst +++ b/docs/compatibility/ml-compatibility/jax-compatibility.rst @@ -47,6 +47,23 @@ with ROCm support: `Community ROCm JAX Docker images `_ 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 ================================================================================