Fix segfault when old GPU plugins are installed.

PiperOrigin-RevId: 726908959
jax-ml · Feb 14, 2025 · 1755c43 · 1755c43
1 parent 4df5961
commit 1755c43
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/jaxlib/gpu_linalg.py b/jaxlib/gpu_linalg.py
@@ -38,12 +38,16 @@
 
 if _cuda_linalg:
   for _name, _value in _cuda_linalg.registrations().items():
+    # TODO(danfm): remove after JAX 0.5.1 release
+    api_version = 1 if "_ffi" in _name else 0
     xla_client.register_custom_call_target(
-        _name, _value, platform="CUDA", api_version=1
+        _name, _value, platform="CUDA", api_version=api_version
     )
 
 if _hip_linalg:
   for _name, _value in _hip_linalg.registrations().items():
+    # TODO(danfm): remove after JAX 0.5.1 release
+    api_version = 1 if "_ffi" in _name else 0
     xla_client.register_custom_call_target(
-        _name, _value, platform="ROCM", api_version=1
+        _name, _value, platform="ROCM", api_version=api_version
     )
diff --git a/jaxlib/gpu_prng.py b/jaxlib/gpu_prng.py
@@ -36,8 +36,10 @@
 
 if _cuda_prng:
   for _name, _value in _cuda_prng.registrations().items():
+    # TODO(danfm): remove after JAX 0.5.1 release
+    api_version = 1 if "_ffi" in _name else 0
     xla_client.register_custom_call_target(_name, _value, platform="CUDA",
-                                           api_version=1)
+                                           api_version=api_version)
 
 for rocm_module_name in [".rocm", "jax_rocm60_plugin"]:
   try:
@@ -51,7 +53,7 @@
 
 if _hip_prng:
   for _name, _value in _hip_prng.registrations().items():
-    # TODO(b/338022728): remove after 6 months, always api_version=1
+    # TODO(danfm): remove after JAX 0.5.1 release
     api_version = 1 if "_ffi" in _name else 0
     xla_client.register_custom_call_target(_name, _value, platform="ROCM",
                                            api_version=api_version)