When importing jax Traceback (most recent call last): File "<stdin>", line 1, in <module> File "/usr/local/lib/python3.7/site-packages/jax/__init__.py", line 22, in <module> from .api import ( File "/usr/local/lib/python3.7/site-packages/jax/api.py", line 37, in <module> from . import core File "/usr/local/lib/python3.7/site-packages/jax/core.py", line 31, in <module> from . import dtypes File "/usr/local/lib/python3.7/site-packages/jax/dtypes.py", line 31, in <module> from .lib import xla_client File "/usr/local/lib/python3.7/site-packages/jax/lib/__init__.py", line 23, in <module> import jaxlib ModuleNotFoundError: No module named 'jaxlib'
*** Bug 254320 has been marked as a duplicate of this bug. ***
jaxlib is a separate library that is also a part of the Jax project on GitHub. Sadly, Google again uses bazel to build it and bazel fails to build it: https://github.com/google/jax/issues/6076
It seems like Bazel can only work inside of Google's environment and fails in the wild.