Windows RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.7.

[huanggou666]

Description

RTX 3060 12GB， cuda 12.1
NVIDIA-SMI 531.68 Driver Version: 531.68 CUDA Version: 12.1 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA GeForce RTX 3060 WDDM | 00000000:01:00.0 On | N/A |
| 0% 44C P5 13W / 170W| 6172MiB / 12288MiB | 2% Default |
| | | N/A |

*bug msg:
File "C:\Users\daihao\AppData\Local\Programs\Python\Python310\lib\site-packages\jax\config.py", line 17, in
from jax.src.config import config # noqa: F401
File "C:\Users\daihao\AppData\Local\Programs\Python\Python310\lib\site-packages\jax_src\config.py", line 24, in
from jax.src import lib
File "C:\Users\daihao\AppData\Local\Programs\Python\Python310\lib\site-packages\jax_src\lib_init.py", line 73, in
version = check_jaxlib_version(
File "C:\Users\daihao\AppData\Local\Programs\Python\Python310\lib\site-packages\jax_src\lib_init.py", line 62, in check_jaxlib_version
raise RuntimeError(msg)
RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.7.
*

What jax/jaxlib version are you using?

jaxlib is version 0.3.25

Which accelerator(s) are you using?

GPU

Additional system info

windows

NVIDIA GPU info

No response

[jakevdp]

Hi - thanks for the question!

The answer is in the error message: you need your jax and jaxlib installations to be compatible. If you want to use jaxlib version 0.3.25, you should use it with jax version 0.3.25. If you want to use jax version 0.4.7, you should use it with jaxlib version 0.4.7. Install matching versions and the error should go away.