AttributeError: _ARRAY_API not found in Numpy 2.0.0 about jax

[superbobry]

I cannot reproduce this in a fresh virtualenv. Can you confirm that doing

python -m venv env
source env/bin/activate
pip install jax==0.4.27 jaxlib==0.4.27 numpy==2.0
python -c "import jax"

reproduces the issue for you?

[superbobry]

It look like the jaxlib package in Conda was compiled for NumPy 1.X? I would recommend to either install JAX from PyPI via pip or open an issue in https://github.com/conda-forge/jaxlib-feedstock.

[hawkinsp]

Closing. We don't provide the conda packages for JAX: please file issues for those against the conda-forge feedstock linked above. The pip packages we ship will work with NumPy 2.0, and we recommend you use those if you want NumPy 2.0 support, even if you're otherwise using conda-forge.

[WhiteSymmetry]

Since my computer only supports the CPU feature, I can install jaxlib's conda and mamba. I tried again today, and it gave the same error.

A module that was compiled using NumPy 1.x cannot be run in
NumPy 2.0.0 as it may crash. To support both 1.x and 2.x
versions of NumPy, modules must be compiled with NumPy 2.0.
Some module may need to rebuild instead e.g. with 'pybind11>=2.12'.

If you are a user of the module, the easiest solution will be to
downgrade to 'numpy<2' or try to upgrade the affected module.
We expect that some modules will need time to support NumPy 2.

Traceback (most recent call last): File "", line 1, in
File "/home//miniconda3/envs/x11/lib/python3.11/site-packages/jax/init.py", line 37, in
import jax.core as _core
File "/home//miniconda3/envs/x11/lib/python3.11/site-packages/jax/core.py", line 18, in
from jax._src.core import (
File "/home//miniconda3/envs/x11/lib/python3.11/site-packages/jax/_src/core.py", line 39, in
from jax._src import dtypes
File "/home//miniconda3/envs/x11/lib/python3.11/site-packages/jax/_src/dtypes.py", line 33, in
from jax._src import config
File "/home//miniconda3/envs/x11/lib/python3.11/site-packages/jax/_src/config.py", line 27, in
from jax._src import lib
File "/home//miniconda3/envs/x11/lib/python3.11/site-packages/jax/_src/lib/init.py", line 87, in
import jaxlib.xla_client as xla_client
File "/home/*/miniconda3/envs/x11/lib/python3.11/site-packages/jaxlib/xla_client.py", line 32, in
from . import xla_extension as _xla
AttributeError: _ARRAY_API not found