Merge pull request jax-ml#10718 from hawkinsp:jaxfix

PiperOrigin-RevId: 448849099
shangshanshi · May 16, 2022 · f572646 · f572646
2 parents 1381afc + 337ec47
commit f572646
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