Merge pull request jax-ml#25084 from ROCm:ci_rocm_version

PiperOrigin-RevId: 700241231

Author: Google-ML-Automation

jax_plugins/rocm/plugin_setup.py

@@ -22,6 +22,11 @@
 project_name = f"jax-rocm{rocm_version}-plugin"
 package_name = f"jax_rocm{rocm_version}_plugin"
 
+# Extract ROCm version from the `ROCM_PATH` environment variable.
+default_rocm_path = "/opt/rocm"
+rocm_path = os.getenv("ROCM_PATH", default_rocm_path)
+rocm_detected_version = rocm_path.split('-')[-1] if '-' in rocm_path else "unknown"
+
 def load_version_module(pkg_path):
   spec = importlib.util.spec_from_file_location(
     'version', os.path.join(pkg_path, 'version.py'))
@@ -43,7 +48,7 @@ def has_ext_modules(self):
     name=project_name,
     version=__version__,
     cmdclass=_cmdclass,
-    description="JAX Plugin for AMD GPUs",
+    description=f"JAX Plugin for AMD GPUs (ROCm:{rocm_detected_version})",
     long_description="",
     long_description_content_type="text/markdown",
     author="Ruturaj4",

jax_plugins/rocm/setup.py

@@ -21,6 +21,11 @@
 project_name = f"jax-rocm{rocm_version}-pjrt"
 package_name = f"jax_plugins.xla_rocm{rocm_version}"
 
+# Extract ROCm version from the `ROCM_PATH` environment variable.
+default_rocm_path = "/opt/rocm"
+rocm_path = os.getenv("ROCM_PATH", default_rocm_path)
+rocm_detected_version = rocm_path.split('-')[-1] if '-' in rocm_path else "unknown"
+
 def load_version_module(pkg_path):
   spec = importlib.util.spec_from_file_location(
     'version', os.path.join(pkg_path, 'version.py'))
@@ -41,7 +46,7 @@ def load_version_module(pkg_path):
 setup(
     name=project_name,
     version=__version__,
-    description="JAX XLA PJRT Plugin for AMD GPUs",
+    description=f"JAX XLA PJRT Plugin for AMD GPUs (ROCm:{rocm_detected_version})",
     long_description="",
     long_description_content_type="text/markdown",
     author="Ruturaj4",