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>
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.