jax report that "No GPU/TPU found, falling back to CPU"

[La-SilverLand]

in the requirement.txt it requires jax>=0.2.4
when incurring this problem as in the title, the jax github homepages says for both support of GPU and CPU,
one needs to install as

pip install --upgrade jax jaxlib==0.1.67+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
The jaxlib version must correspond to the version of the existing CUDA installation you want to use

but that jaxlib repo link does not contain any version >= 0.2.4,
actually the newest version there is 0.1.67

How do you install the environment ?

[La-SilverLand]

I got it fixed, the jax>=0.2.4 does not mean jaxlib>=0.2.4
for my situation, i have cuda 10.2 and jax==0.2.13
jaxlib==0.1.65+cuda102 is suitable

[DaShenZi721]

@La-SilverLand Hello! Sorry to bother you. Have you ever encountered the following problem? I think it may be related to the version of flax.

Traceback (most recent call last):
  File "lra_benchmarks/listops/train.py", line 28, in <module>
    from flax.deprecated import nn
ModuleNotFoundError: No module named 'flax.deprecated'

[DaShenZi721]

@La-SilverLand This is my setting:

• cuda: 11.2
• jax: 0.2.13
• jaxlib: 0.1.65+cuda112
• flax: 0.2.2
• tensorflow: 2.7.0