Default installation of JAX + Tensorflow produces inconsistent environment

[joezuntz]

Solution to issue cannot be found in the documentation.

• I checked the documentation.

Issue

Installing default versions of both JAX and Tensorflow from conda-forge currently results in an inconsistent environment in which neither of them can be imported. Using either of the two separately is fine.

I can recreate this in the miniforge3 docker image like this

FROM condaforge/miniforge3
RUN conda install tensorflow jax
RUN conda list
RUN conda info
RUN python -c 'import tensorflow'
# this also fails
# RUN python -c 'import jax'

which gives the following error:

#7 [5/5] RUN python -c 'import tensorflow'
#7 sha256:c7f1bae8ee24e9614e6e17d5be2614442a2df455b7b3733c3f1e4b8b7a40132a
#7 2.331 Traceback (most recent call last):
#7 2.331   File "<string>", line 1, in <module>
#7 2.331   File "/opt/conda/lib/python3.9/site-packages/tensorflow/__init__.py", line 55, in <module>
#7 2.331     from ._api.v2 import compat
#7 2.331   File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/__init__.py", line 39, in <module>
#7 2.331     from . import v1
#7 2.331   File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/__init__.py", line 34, in <module>
#7 2.331     from . import compat
#7 2.331   File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/__init__.py", line 39, in <module>
#7 2.333     from . import v1
#7 2.333   File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/v1/__init__.py", line 51, in <module>
#7 2.334     from tensorflow._api.v2.compat.v1 import lite
#7 2.334   File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/__init__.py", line 11, in <module>
#7 2.334     from . import experimental
#7 2.334   File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/__init__.py", line 10, in <module>
#7 2.334     from . import authoring
#7 2.334   File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/authoring/__init__.py", line 10, in <module>
#7 2.334     from tensorflow.lite.python.authoring.authoring import compatible
#7 2.334   File "/opt/conda/lib/python3.9/site-packages/tensorflow/lite/python/authoring/authoring.py", line 43, in <module>
#7 2.335     from tensorflow.lite.python import convert
#7 2.335   File "/opt/conda/lib/python3.9/site-packages/tensorflow/lite/python/convert.py", line 33, in <module>
#7 2.335     from tensorflow.lite.python import util
#7 2.335   File "/opt/conda/lib/python3.9/site-packages/tensorflow/lite/python/util.py", line 55, in <module>
#7 2.335     from jax import xla_computation as _xla_computation
#7 2.335   File "/opt/conda/lib/python3.9/site-packages/jax/__init__.py", line 37, in <module>
#7 2.335     from jax import config as _config_module
#7 2.335   File "/opt/conda/lib/python3.9/site-packages/jax/config.py", line 18, in <module>
#7 2.335     from jax._src.config import config
#7 2.335   File "/opt/conda/lib/python3.9/site-packages/jax/_src/config.py", line 27, in <module>
#7 2.335     from jax._src import lib
#7 2.335   File "/opt/conda/lib/python3.9/site-packages/jax/_src/lib/__init__.py", line 101, in <module>
#7 2.335     version = check_jaxlib_version(
#7 2.335   File "/opt/conda/lib/python3.9/site-packages/jax/_src/lib/__init__.py", line 90, in check_jaxlib_version
#7 2.336     raise RuntimeError(msg)
#7 2.336 RuntimeError: jaxlib is version 0.1.75, but this version of jax requires version >= 0.3.2.
#7 ERROR: executor failed running [/bin/sh -c python -c 'import tensorflow']: exit code: 1

Installed packages

#5 [3/5] RUN conda list
#5 sha256:1b4a2bef8e285ea7908ee767af6ddccd2b54750f9f686f79b9cdf637ae8dbacc
#5 1.412 # packages in environment at /opt/conda:
#5 1.412 #
#5 1.412 # Name                    Version                   Build  Channel
#5 1.412 _libgcc_mutex             0.1                 conda_forge    conda-forge
#5 1.412 _openmp_mutex             4.5                       1_gnu    conda-forge
#5 1.412 abseil-cpp                20210324.2           h9c3ff4c_0    conda-forge
#5 1.412 absl-py                   1.0.0              pyhd8ed1ab_0    conda-forge
#5 1.412 aiohttp                   3.8.1            py39hb9d737c_1    conda-forge
#5 1.412 aiosignal                 1.2.0              pyhd8ed1ab_0    conda-forge
#5 1.412 astunparse                1.6.3              pyhd8ed1ab_0    conda-forge
#5 1.412 async-timeout             4.0.2              pyhd8ed1ab_0    conda-forge
#5 1.412 attrs                     21.4.0             pyhd8ed1ab_0    conda-forge
#5 1.412 blinker                   1.4                        py_1    conda-forge
#5 1.412 brotlipy                  0.7.0           py39h3811e60_1003    conda-forge
#5 1.412 bzip2                     1.0.8                h7f98852_4    conda-forge
#5 1.412 c-ares                    1.18.1               h7f98852_0    conda-forge
#5 1.412 ca-certificates           2021.10.8            ha878542_0    conda-forge
#5 1.412 cached-property           1.5.2                hd8ed1ab_1    conda-forge
#5 1.412 cached_property           1.5.2              pyha770c72_1    conda-forge
#5 1.412 cachetools                4.2.4              pyhd8ed1ab_0    conda-forge
#5 1.412 certifi                   2021.10.8        py39hf3d152e_2    conda-forge
#5 1.412 cffi                      1.15.0           py39h4bc2ebd_0    conda-forge
#5 1.412 charset-normalizer        2.0.12             pyhd8ed1ab_0    conda-forge
#5 1.412 click                     8.1.3            py39hf3d152e_0    conda-forge
#5 1.412 colorama                  0.4.4              pyh9f0ad1d_0    conda-forge
#5 1.412 conda                     4.12.0           py39hf3d152e_0    conda-forge
#5 1.412 conda-package-handling    1.8.0            py39hb9d737c_0    conda-forge
#5 1.412 cryptography              36.0.2           py39hd97740a_0    conda-forge
#5 1.412 frozenlist                1.3.0            py39hb9d737c_1    conda-forge
#5 1.412 gast                      0.4.0              pyh9f0ad1d_0    conda-forge
#5 1.412 giflib                    5.2.1                h36c2ea0_2    conda-forge
#5 1.412 google-auth               1.35.0             pyh6c4a22f_0    conda-forge
#5 1.412 google-auth-oauthlib      0.4.6              pyhd8ed1ab_0    conda-forge
#5 1.412 google-pasta              0.2.0              pyh8c360ce_0    conda-forge
#5 1.412 grpc-cpp                  1.42.0               ha1441d3_1    conda-forge
#5 1.412 grpcio                    1.42.0           py39hff7568b_0    conda-forge
#5 1.412 h5py                      3.6.0           nompi_py39h7e08c79_100    conda-forge
#5 1.412 hdf5                      1.12.1          nompi_h2386368_104    conda-forge
#5 1.412 icu                       69.1                 h9c3ff4c_0    conda-forge
#5 1.412 idna                      3.3                pyhd8ed1ab_0    conda-forge
#5 1.412 importlib-metadata        4.11.3           py39hf3d152e_1    conda-forge
#5 1.412 jax                       0.3.7              pyhd8ed1ab_0    conda-forge
#5 1.412 jaxlib                    0.1.75           py39hde0f152_0    conda-forge
#5 1.412 jpeg                      9e                   h166bdaf_1    conda-forge
#5 1.412 keras                     2.7.0              pyhd8ed1ab_0    conda-forge
#5 1.412 keras-preprocessing       1.1.2              pyhd8ed1ab_0    conda-forge
#5 1.412 keyutils                  1.6.1                h166bdaf_0    conda-forge
#5 1.412 krb5                      1.19.3               h3790be6_0    conda-forge
#5 1.412 ld_impl_linux-64          2.36.1               hea4e1c9_2    conda-forge
#5 1.412 libblas                   3.9.0           14_linux64_openblas    conda-forge
#5 1.412 libcblas                  3.9.0           14_linux64_openblas    conda-forge
#5 1.412 libcurl                   7.83.0               h7bff187_0    conda-forge
#5 1.412 libedit                   3.1.20191231         he28a2e2_2    conda-forge
#5 1.412 libev                     4.33                 h516909a_1    conda-forge
#5 1.412 libffi                    3.4.2                h7f98852_5    conda-forge
#5 1.412 libgcc-ng                 11.2.0              h1d223b6_14    conda-forge
#5 1.412 libgfortran-ng            11.2.0              h69a702a_16    conda-forge
#5 1.412 libgfortran5              11.2.0              h5c6108e_16    conda-forge
#5 1.412 libgomp                   11.2.0              h1d223b6_14    conda-forge
#5 1.412 liblapack                 3.9.0           14_linux64_openblas    conda-forge
#5 1.412 libnghttp2                1.47.0               h727a467_0    conda-forge
#5 1.412 libnsl                    2.0.0                h7f98852_0    conda-forge
#5 1.412 libopenblas               0.3.20          pthreads_h78a6416_0    conda-forge
#5 1.412 libpng                    1.6.37               h21135ba_2    conda-forge
#5 1.412 libprotobuf               3.19.4               h780b84a_0    conda-forge
#5 1.412 libssh2                   1.10.0               ha56f1ee_2    conda-forge
#5 1.412 libstdcxx-ng              11.2.0              he4da1e4_16    conda-forge
#5 1.412 libuuid                   2.32.1            h7f98852_1000    conda-forge
#5 1.412 libzlib                   1.2.11            h36c2ea0_1013    conda-forge
#5 1.412 markdown                  3.3.7              pyhd8ed1ab_0    conda-forge
#5 1.412 multidict                 6.0.2            py39hb9d737c_1    conda-forge
#5 1.412 ncurses                   6.3                  h9c3ff4c_0    conda-forge
#5 1.412 numpy                     1.22.3           py39hc58783e_2    conda-forge
#5 1.412 oauthlib                  3.2.0              pyhd8ed1ab_0    conda-forge
#5 1.412 openssl                   1.1.1o               h166bdaf_0    conda-forge
#5 1.412 opt_einsum                3.3.0              pyhd8ed1ab_1    conda-forge
#5 1.412 pip                       22.0.4             pyhd8ed1ab_0    conda-forge
#5 1.412 protobuf                  3.19.4           py39he80948d_0    conda-forge
#5 1.412 pyasn1                    0.4.8                      py_0    conda-forge
#5 1.412 pyasn1-modules            0.2.7                      py_0    conda-forge
#5 1.412 pycosat                   0.6.3           py39h3811e60_1009    conda-forge
#5 1.412 pycparser                 2.21               pyhd8ed1ab_0    conda-forge
#5 1.412 pyjwt                     2.3.0              pyhd8ed1ab_1    conda-forge
#5 1.412 pyopenssl                 22.0.0             pyhd8ed1ab_0    conda-forge
#5 1.412 pysocks                   1.7.1            py39hf3d152e_4    conda-forge
#5 1.412 python                    3.9.10          h85951f9_2_cpython    conda-forge
#5 1.412 python-flatbuffers        2.0                pyhd8ed1ab_0    conda-forge
#5 1.412 python_abi                3.9                      2_cp39    conda-forge
#5 1.412 pyu2f                     0.1.5              pyhd8ed1ab_0    conda-forge
#5 1.412 re2                       2021.11.01           h9c3ff4c_0    conda-forge
#5 1.412 readline                  8.1                  h46c0cb4_0    conda-forge
#5 1.412 requests                  2.27.1             pyhd8ed1ab_0    conda-forge
#5 1.412 requests-oauthlib         1.3.1              pyhd8ed1ab_0    conda-forge
#5 1.412 rsa                       4.8                pyhd8ed1ab_0    conda-forge
#5 1.412 ruamel_yaml               0.15.80         py39h3811e60_1006    conda-forge
#5 1.412 scipy                     1.8.0            py39hee8e79c_1    conda-forge
#5 1.412 setuptools                60.10.0          py39hf3d152e_0    conda-forge
#5 1.412 six                       1.16.0             pyh6c4a22f_0    conda-forge
#5 1.412 snappy                    1.1.8                he1b5a44_3    conda-forge
#5 1.412 sqlite                    3.37.1               h4ff8645_0    conda-forge
#5 1.412 tensorboard               2.6.0              pyhd8ed1ab_1    conda-forge
#5 1.412 tensorboard-data-server   0.6.0            py39hd97740a_2    conda-forge
#5 1.412 tensorboard-plugin-wit    1.8.1              pyhd8ed1ab_0    conda-forge
#5 1.412 tensorflow                2.7.0           cpu_py39h4655687_0    conda-forge
#5 1.412 tensorflow-base           2.7.0           cpu_py39hf4995fd_0    conda-forge
#5 1.412 tensorflow-estimator      2.7.0           cpu_py39ha241409_0    conda-forge
#5 1.412 termcolor                 1.1.0                      py_2    conda-forge
#5 1.412 tk                        8.6.12               h27826a3_0    conda-forge
#5 1.412 tqdm                      4.63.0             pyhd8ed1ab_0    conda-forge
#5 1.412 typing-extensions         4.2.0                hd8ed1ab_1    conda-forge
#5 1.412 typing_extensions         4.2.0              pyha770c72_1    conda-forge
#5 1.412 tzdata                    2022a                h191b570_0    conda-forge
#5 1.412 urllib3                   1.26.9             pyhd8ed1ab_0    conda-forge
#5 1.412 werkzeug                  2.1.2              pyhd8ed1ab_1    conda-forge
#5 1.412 wheel                     0.37.1             pyhd8ed1ab_0    conda-forge
#5 1.412 wrapt                     1.14.1           py39hb9d737c_0    conda-forge
#5 1.412 xz                        5.2.5                h516909a_1    conda-forge
#5 1.412 yaml                      0.2.5                h7f98852_2    conda-forge
#5 1.412 yarl                      1.7.2            py39hb9d737c_2    conda-forge
#5 1.412 zipp                      3.8.0              pyhd8ed1ab_0    conda-forge
#5 1.412 zlib                      1.2.11            h36c2ea0_1013    conda-forge

Environment info

#6 [4/5] RUN conda info
#6 sha256:e04dec7bc9e71b96d819465250e444e4080f1874b5342685a25c0744c32e910d
#6 0.730 
#6 0.730      active environment : None
#6 0.730        user config file : /root/.condarc
#6 0.730  populated config files : /opt/conda/.condarc
#6 0.730           conda version : 4.12.0
#6 0.730     conda-build version : not installed
#6 0.730          python version : 3.9.10.final.0
#6 0.730        virtual packages : __linux=5.10.47=0
#6 0.730                           __glibc=2.31=0
#6 0.730                           __unix=0=0
#6 0.730                           __archspec=1=x86_64
#6 0.730        base environment : /opt/conda  (writable)
#6 0.730       conda av data dir : /opt/conda/etc/conda
#6 0.730   conda av metadata url : None
#6 0.730            channel URLs : https://conda.anaconda.org/conda-forge/linux-64
#6 0.730                           https://conda.anaconda.org/conda-forge/noarch
#6 0.730           package cache : /opt/conda/pkgs
#6 0.730                           /root/.conda/pkgs
#6 0.730        envs directories : /opt/conda/envs
#6 0.730                           /root/.conda/envs
#6 0.730                platform : linux-64
#6 0.730              user-agent : conda/4.12.0 requests/2.27.1 CPython/3.9.10 Linux/5.10.47-linuxkit ubuntu/20.04.4 glibc/2.31
#6 0.730                 UID:GID : 0:0
#6 0.730              netrc file : None
#6 0.730            offline mode : False
#6 0.730

[hmaarrfk]

I can recreate. Thank you for the detailed information.

I think that:

1. The current jax feedstock should add constraints to the jaxlib dependency and write a repopatch to patch existing builds. cc @conda-forge/jax @conda-forge/jaxlib
2. We are aware that the tensorflow feedstock has become incompatible with the latest packages from conda-forge. A few PRs are trying to alleviate that, but it might require more work due to abseil library mismatches.

If you could:

1. Try to manually specify the versions you want. You will have to downgrade the version of jax you want is my guess.

mamba install "tensorflow>=2.7.0" "jax=x.x.x" "jaxlib=y.y.y."

that would help alleviate things
2. Try to submit repo-data patches to require the appropriate versions of jaxlib for jax.
3. Submit a pull request to jax fixing the dependnecy for future packages.

[ngam]

I believe this should be resolved now, no? At least the current packages are compatible. A repodata patch may be worthy tho

[joezuntz]

If I add RUN conda update conda to the Dockerfile I described above then this is indeed fixed.

If I don't do this then the default base image, condaforge/miniforge3:latest is still inconsistent.

[ngam]

If I add RUN conda update conda to the Dockerfile I described above then this is indeed fixed.

If I don't do this then the default base image, condaforge/miniforge3:latest is still inconsistent.

Okay, I see.

Btw, is this for a GPU or CPU deployment? I highly recommend you follow hmaarrfk's recommendation above. Our jaxlib>=0.3.10 (thus jax>=0.3.10) and tensorflow>=2.8.1 are somewhat different from previous versions (enabling XLA for tf and adding cuda for jax). Also, both are slightly more advanced in terms of migrations (they're incompatible with things like gdal, etc. as the latter haven't been migrated to the latest icu).

[joezuntz]

I've only tested on CPU.

I was indeed able to work around things with @hmaarrfk's suggestions above, so thanks for that!

[ngam]

Who is in charge of condaforge/miniforge3:latest? We can just bump that up?

Re CPU vs GPU, I am asking because our current recommendation is that users specify *cuda* explicitly, e.g. tensorflow>=2.8.1=*cuda* and jaxlib>=0.3.10=*cuda*. In case you need that info in the future :)