AttributeError: module 'jaxlib.xla_extension' has no attribute 'PmapFunction'

[PRISHIta123]

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.

[PRISHIta123]

@yashk2810 please check your latest commit

[yashk2810]

Which version were you running with previously?

Any latest commit would not affect you since you are installing an already built version.

[PRISHIta123]

Which version were you running with previously?

Any latest commit would not affect you since you are installing an already built version.

This fetches the latest version right?

!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

How do I specify the previous version while installing?

[PRISHIta123]

I downgraded using this which I found on StackOverflow:

!pip install --upgrade jax==0.3.10 jaxlib==0.3.10+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_releases.html

It would be helpful if this could be mentioned on the README page as it becomes difficult to revert back to older versions with only the general installation command that is provided, when newer versions are released and some functionalities suddenly do not work.

[yashk2810]

We'll release with the fix again! Sorry for the breakage.

[PRISHIta123]

No worries!

[yashk2810]

We released jax 0.3.12: https://pypi.org/project/jax/0.3.12/ with the fix.

[PRISHIta123]

Ok, thanks for addressing this.