AttributeError: module 'jaxlib.xla_extension' has no attribute 'PmapFunction'
I am getting the above error after installing the latest version of jax and trying to import it for some computations:
```
!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
```
```
/usr/local/lib/python3.7/dist-packages/jax/_src/lax/linalg.py in <module>()
    977       lu_pivots_to_permutation_p,
    978       partial(_lu_pivots_to_permutation_gpu_lowering,
--> 979               gpu_linalg.cuda_lu_pivots_to_permutation),
    980       platform='cuda')
    981   mlir.register_lowering(

AttributeError: module 'jaxlib.cuda_linalg' has no attribute 'cuda_lu_pivots_to_permutation'
```

It would be helpful if someone could please look into this ASAP. This is a new error I received when running my code within the past hour and had no problem with this earlier. 
diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
index ea9b2d9a8..f45bf7295 100644
--- a/jax/_src/lib/__init__.py
+++ b/jax/_src/lib/__init__.py
@@ -183,7 +183,7 @@ except ImportError:
   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
 
['jax/_src/lib/__init__.py']