Install with GPU support

[seanjkanderson]

I'm having trouble getting GPJax to install correctly on an Ubuntu server. I'm using the server because it has GPU support, and in general I can install jax and it's dependencies and run a script (without GPJax) with GPU support.

So I took a look at the installation page for GPJax: https://gpjax.readthedocs.io/en/latest/installation.html
and also tried to follow the Jax install guide: https://github.com/google/jax#installation

If I install GPJax first, then when I go to setup CUDA support with
pip install "jax[cuda11_cudnn805]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
I get

Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Requirement already satisfied: jax[cuda11_cudnn805] in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (0.3.16)
Requirement already satisfied: opt-einsum in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from jax[cuda11_cudnn805]) (3.3.0)
Requirement already satisfied: etils[epath] in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from jax[cuda11_cudnn805]) (0.7.1)
Requirement already satisfied: typing-extensions in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from jax[cuda11_cudnn805]) (4.3.0)
Requirement already satisfied: scipy>=1.5 in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from jax[cuda11_cudnn805]) (1.9.0)
Requirement already satisfied: absl-py in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from jax[cuda11_cudnn805]) (1.2.0)
Requirement already satisfied: numpy>=1.20 in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from jax[cuda11_cudnn805]) (1.23.2)
Collecting jaxlib==0.3.15+cuda11.cudnn805
  Downloading https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.15%2Bcuda11.cudnn805-cp39-none-manylinux2014_x86_64.whl (250.6 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 250.6/250.6 MB 7.5 MB/s eta 0:00:00
Requirement already satisfied: zipp in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from etils[epath]->jax[cuda11_cudnn805]) (3.8.1)
Requirement already satisfied: importlib_resources in /home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages (from etils[epath]->jax[cuda11_cudnn805]) (5.9.0)
Installing collected packages: jaxlib
  Attempting uninstall: jaxlib
    Found existing installation: jaxlib 0.3.15+cuda11.cudnn82
    Uninstalling jaxlib-0.3.15+cuda11.cudnn82:
      Successfully uninstalled jaxlib-0.3.15+cuda11.cudnn82
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
gpjax 0.4.10 requires jax==0.3.5, but you have jax 0.3.16 which is incompatible.
gpjax 0.4.10 requires jaxlib==0.3.5, but you have jaxlib 0.3.15+cuda11.cudnn805 which is incompatible.
Successfully installed jaxlib-0.3.15+cuda11.cudnn805

Now if I try to naively run my scripts that utilize GPJax, I get the following error:

/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/chex/_src/pytypes.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
  PyTreeDef = type(jax.tree_structure(None))
Precomputed reusables
Created GPExperiment instance
Traceback (most recent call last):
  File "/home/paperspace/pycharm_projects/exp_design/basic_experiment.py", line 78, in <module>
    gp_exp.train_model(dataset)
  File "/home/paperspace/pycharm_projects/exp_design/gp_learn.py", line 235, in train_model
    params = gpx.transform(params, unconstrainer)
  File "/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/gpjax/parameters.py", line 147, in transform
    return jax.tree_map(lambda param, trans: trans(param), params, transform_map)
  File "/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/jax/_src/tree_util.py", line 205, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/jax/_src/tree_util.py", line 205, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/gpjax/parameters.py", line 147, in <lambda>
    return jax.tree_map(lambda param, trans: trans(param), params, transform_map)
  File "/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/distrax/_src/bijectors/lambda_bijector.py", line 112, in inverse
    return self._inverse(y)
  File "/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/distrax/_src/utils/transformations.py", line 119, in wrapped
    out = _interpret_inverse(jaxpr, consts, *args)
  File "/home/paperspace/anaconda3/envs/py39/lib/python3.9/site-packages/distrax/_src/utils/transformations.py", line 250, in _interpret_inverse
    write(jax.core.unitvar, jax.core.unit)
AttributeError: module 'jax.core' has no attribute 'unitvar'

I'm running this on Ubuntu 20.04.3 and have tried python 3.8-3.10 without luck. If I don't use the GPU support, I can run my scripts but obviously don't get the GPU benefits.

I would similarly be interested in knowing how to install on a colab notebook as !pip install gpjax results in:

ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
thinc 8.1.0 requires typing-extensions<4.2.0,>=3.7.4.1; python_version < "3.8", but you have typing-extensions 4.3.0 which is incompatible.
spacy 3.4.1 requires typing-extensions<4.2.0,>=3.7.4; python_version < "3.8", but you have typing-extensions 4.3.0 which is incompatible.
Successfully installed chex-0.1.3 distrax-0.1.2 gpjax-0.4.10 jax-0.3.5 jaxlib-0.3.5 jaxtyping-0.1.0 keras-2.8.0 ml-collections-0.1.0 optax-0.1.3 protobuf-3.19.0 tensorboard-2.8.0 tensorflow-2.8.1+zzzcolab20220518083849 tensorflow-estimator-2.8.0 tensorflow-probability-0.16.0 typeguard-2.13.3 typing-extensions-4.3.0
WARNING: The following packages were previously imported in this runtime:
  [google]
You must restart the runtime in order to use newly installed versions.

Thanks!

[seanjkanderson]

Moving to Issues

[maxhinne]

Have you perhaps found a workaround for this issue (I can't find it under Issues, but it does not seem to be fixed yet - I encounter the same problem using Google Colab)?

[seanjkanderson]

Yes, this should do it: #102 (comment)