Missing dependency: jaxlib/ pip version conflict

[mollenhauerm]

Installing via

python setup.py install

gives

Python 3.7.3 (default, Mar 27 2019, 22:11:17) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jaxrk
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/mattes/projects/JaxRK/jaxrk/__init__.py", line 1, in <module>
    from . import (rkhs, kern)
  File "/home/mattes/projects/JaxRK/jaxrk/rkhs/__init__.py", line 2, in <module>
    from .vector import FiniteVec, inner, CombVec
  File "/home/mattes/projects/JaxRK/jaxrk/rkhs/vector.py", line 2, in <module>
    import jax.numpy as np
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/__init__.py", line 19, in <module>
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/api.py", line 37, in <module>
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/core.py", line 30, in <module>
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/dtypes.py", line 31, in <module>
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/lib/__init__.py", line 18, in <module>
ModuleNotFoundError: No module named 'jaxlib'

Manually adding jaxlib to the dependencies in setup.py gives

Python 3.7.3 (default, Mar 27 2019, 22:11:17) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jaxrk
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/mattes/projects/JaxRK/jaxrk/__init__.py", line 1, in <module>
    from . import (rkhs, kern)
  File "/home/mattes/projects/JaxRK/jaxrk/rkhs/__init__.py", line 2, in <module>
    from .vector import FiniteVec, inner, CombVec
  File "/home/mattes/projects/JaxRK/jaxrk/rkhs/vector.py", line 2, in <module>
    import jax.numpy as np
  File "<frozen importlib._bootstrap>", line 983, in _find_and_load
  File "<frozen importlib._bootstrap>", line 967, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 668, in _load_unlocked
  File "<frozen importlib._bootstrap>", line 638, in _load_backward_compatible
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/__init__.py", line 19, in <module>
  File "<frozen importlib._bootstrap>", line 983, in _find_and_load
  File "<frozen importlib._bootstrap>", line 967, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 668, in _load_unlocked
  File "<frozen importlib._bootstrap>", line 638, in _load_backward_compatible
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/api.py", line 37, in <module>
  File "<frozen importlib._bootstrap>", line 983, in _find_and_load
  File "<frozen importlib._bootstrap>", line 967, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 668, in _load_unlocked
  File "<frozen importlib._bootstrap>", line 638, in _load_backward_compatible
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/core.py", line 30, in <module>
  File "<frozen importlib._bootstrap>", line 983, in _find_and_load
  File "<frozen importlib._bootstrap>", line 967, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 668, in _load_unlocked
  File "<frozen importlib._bootstrap>", line 638, in _load_backward_compatible
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/dtypes.py", line 31, in <module>
  File "<frozen importlib._bootstrap>", line 983, in _find_and_load
  File "<frozen importlib._bootstrap>", line 967, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 668, in _load_unlocked
  File "<frozen importlib._bootstrap>", line 638, in _load_backward_compatible
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/lib/__init__.py", line 45, in <module>
  File "/home/mattes/anaconda3/envs/jaxrk-test/lib/python3.7/site-packages/jax-0.1.62-py3.7.egg/jax/lib/__init__.py", line 43, in _check_jaxlib_version
ValueError: jaxlib is version 0.1.23, but this version of jax requires version 0.1.41.

A common cause of this error is that you installed jaxlib using pip, but your version of pip is too old to support manylinux2010 wheels. Try running:

pip install --upgrade pip
pip install --upgrade jax jaxlib

with an up-to-date pip (20.0.2).

running

pip install --upgrade jax jaxlib

solves the issue, but is somehow tedious...

[mollenhauerm]

potentially relevant:
jax-ml/jax#1792

[ingmarschuster]

This indeed seems to be a jax installation issue - it seems jax should add jaxlib as a requirement. In a Linux environment I checked this with conda:

conda  create --name jaxrk
conda activate jaxrk
pip install --upgrade pip jaxlib jax jaxrk

Which enabled running the quick start in a clean conda environment without problems. Can you confirm? The environment can then be deleted again using

conda deactivate
conda env remove --name jaxrk

[mollenhauerm]

pip install --upgrade pip jaxlib jax jaxrk seems to do the trick. The example notebook is running fine now.