Fix jax 0.3.11 GPU breakge when used with jaxlib 0.3.10.

google · May 16, 2022 · 337ec47 · 337ec47
1 parent 1381afc
commit 337ec47
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
@@ -183,7 +183,7 @@ def _parse_version(v: str) -> Tuple[int, ...]:
   hip_linalg = None
 
 try:
-  import jaxlib.cuda_linalg as gpu_linalg  # pytype: disable=import-error
+  import jaxlib.gpu_linalg as gpu_linalg  # pytype: disable=import-error
 except ImportError:
   gpu_linalg = None