ImportError: cannot import name 'cusolver' from 'jaxlib'

[nprime496]

Hi everyone, I'm not able to import cusolver in any installable (>=0.1.60) version of jaxlib. I was working for version 0.1.55 which seems not to be available anymore.

from jaxlib import cusolver

[jakevdp]

Can you say more about the system you're using? This seems to work, for example, on a Colab CPU or GPU runtime with jaxlib 0.1.70:

import jaxlib
print(jaxlib.__version__)
# 0.1.70
from jaxlib import cusolver
print(cusolver)
# <module 'jaxlib.cusolver' from '/usr/local/lib/python3.7/dist-packages/jaxlib/cusolver.py'>

Also, note that if you want to install older jaxlib versions, while they're not available on PyPI, you can use a command like the following:

!pip install jaxlib==0.1.55 -f https://storage.googleapis.com/jax-releases/jax_releases.html 

(we have to periodically remove older versions from PyPI because of their storage limits)

[nprime496]

Your are right! This code works well on Colab. Actually I'm having the error in Github Actions env, I don't know why but the same code fails at from jaxlib import cusolver part.