diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
index ea9b2d9a844f..f45bf729551d 100644
--- 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