module 'jax' has no attribute 'api'

[khanhpham3799]

when calling: state = jax_utils.replicate(state, devices=devices)
the terminal return: AttributeError: module 'jax' has no attribute 'api'

[xh-liu-tech]

Same issue.
Upgrading flax from 0.3.4 to 0.3.5 solves the problem for me.

[khanhpham3799]

i have tried but it said that the latest version only up to 0.3.4

[xh-liu-tech]

i have tried but it said that the latest version only up to 0.3.4

Maybe you need to upgrade pip3 first.
You can also try to install it offline using a whl file.

[Tomzxcv]

Same issue. Upgrading flax from 0.3.4 to 0.3.5 solves the problem for me.

I get another error when upgrading flax to 0.3.5
Traceback (most recent call last):
File "/home/lhm/NeRF/nerfies/train.py", line 23, in
from flax import jax_utils
File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/init.py", line 37, in
from . import linen
File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/init.py", line 23, in
from .attention import (MultiHeadDotProductAttention, SelfAttention,
File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/attention.py", line 26, in
from flax.linen.linear import default_kernel_init
File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/linear.py", line 21, in
from flax.linen.module import Module, compact
File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/module.py", line 31, in
from jax._src.numpy.lax_numpy import isin
ImportError: cannot import name 'isin' from 'jax._src.numpy.lax_numpy' (/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py)

[Tomzxcv]

Same issue. Upgrading flax from 0.3.4 to 0.3.5 solves the problem for me.

I get another error when upgrading flax to 0.3.5 Traceback (most recent call last): File "/home/lhm/NeRF/nerfies/train.py", line 23, in from flax import jax_utils File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/init.py", line 37, in from . import linen File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/init.py", line 23, in from .attention import (MultiHeadDotProductAttention, SelfAttention, File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/attention.py", line 26, in from flax.linen.linear import default_kernel_init File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/linear.py", line 21, in from flax.linen.module import Module, compact File "/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/flax/linen/module.py", line 31, in from jax._src.numpy.lax_numpy import isin ImportError: cannot import name 'isin' from 'jax._src.numpy.lax_numpy' (/home/lhm/anaconda3/envs/nerf-env/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py)

solver by : "pip install jax==0.3.4 jaxlib==0.3.2"