AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'

Please:

• Check for duplicate issues.
• Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

Code:

import jax
import jaxlib
import jaxlib.xla_extension as xe

jax.random.PRNGKey(0)
print(jax.devices())
print(xe.CpuDevice)

• If applicable, include full error messages/tracebacks.

Output:

[GpuDevice(id=0, process_index=0)]
Traceback (most recent call last):
  File "~/jax/1_installation_test.py", line 8, in <module>
    print(xe.CpuDevice)
AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'

Environment details:

• Fedora 35 with cuda 11.6, cudnn 8.4
• jax 0.3.7
• jaxlib 0.3.7+cuda11.cudnn82
• python 3.9

[mblondel]

I ran into the same issue and solved it by pip install chex --upgrade.

[ozencgungor]

I just ran into the exact same error.

Environment details:

cuda 11.2, cudnn 8.1
jax 0.3.7
jaxlib 0.3.7+cuda11.cudnn82
python 3.8.6

Just updated from jax/jaxlib 0.3.6 to 0.3.7, took care to install the right cuda and cudnn version. I’m guessing it has something to do with v0.3.7.

Hoping it gets solved quickly.

Update: Just downgraded to jax 0.3.6 and jaxlib 0.3.5+cuda11.cudnn805 and it seems to work for now.