jax version error

[Prindle19]

Imports section fails:

RuntimeError Traceback (most recent call last) in <cell line: 12>() 10 import cartopy.crs as ccrs 11 from google.cloud import storage ---> 12 from graphcast import autoregressive 13 from graphcast import casting 14 from graphcast import checkpoint

7 frames /usr/local/lib/python3.10/dist-packages/jax/_src/lib/init.py in check_jaxlib_version(jax_version, jaxlib_version, minimum_jaxlib_version) 61 msg = (f'jaxlib is version {jaxlib_version}, but this version ' 62 f'of jax requires version >= {minimum_jaxlib_version}.') ---> 63 raise RuntimeError(msg) 64 65 if _jaxlib_version > _jax_version:

RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.14.

[scploeger]

I had this issue too. I tried updgrading jax to a version >= 0.4.14 with !pip install jaxlib==0.4.14, but this line from graphcast import casting still causes the following error:

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-3-08ed8f8071bd> in <cell line: 13>()
     11 from google.cloud import storage
     12 from graphcast import autoregressive
---> 13 from graphcast import casting
     14 from graphcast import checkpoint
     15 from graphcast import data_utils

4 frames

/usr/local/lib/python3.10/dist-packages/graphcast/casting.py in <module>
     17 from typing import Any, Mapping, Tuple
     18 
---> 19 import chex
     20 from graphcast import predictor_base
     21 import haiku as hk

/usr/local/lib/python3.10/dist-packages/chex/__init__.py in <module>
     15 """Chex: Testing made fun, in JAX!"""
     16 
---> 17 from chex._src.asserts import assert_axis_dimension
     18 from chex._src.asserts import assert_axis_dimension_comparator
     19 from chex._src.asserts import assert_axis_dimension_gt

/usr/local/lib/python3.10/dist-packages/chex/_src/asserts.py in <module>
     24 from unittest import mock
     25 
---> 26 from chex._src import asserts_internal as _ai
     27 from chex._src import pytypes
     28 import jax

/usr/local/lib/python3.10/dist-packages/chex/_src/asserts_internal.py in <module>
     32 
     33 from absl import logging
---> 34 from chex._src import pytypes
     35 import jax
     36 from jax.experimental import checkify

/usr/local/lib/python3.10/dist-packages/chex/_src/pytypes.py in <module>
     23 ArrayBatched = jax.interpreters.batching.BatchTracer
     24 ArrayNumpy = np.ndarray
---> 25 ArraySharded = jax.interpreters.pxla.ShardedDeviceArray
     26 # For instance checking, use `isinstance(x, jax.Array)`.
     27 if hasattr(jax, 'Array'):

AttributeError: module 'jax.interpreters.pxla' has no attribute 'ShardedDeviceArray'

[scploeger]

I was above to temporarily get past this issue by downgrading the following packages:

!pip install jax==0.4.16
!pip install jaxlib==0.4.16
!pip install flax==0.7.3

[welcomethebird]

It looks like there is a runtime error related to the JAX library. The error message indicates that the version of JAX you have is not compatible with the version required by your code. Quickly Check Google's cloud library version & JAX version requirements.

[scploeger]

!pip install jax==0.4.16 !pip install jaxlib==0.4.16 !pip install flax==0.7.3

@Shadowwomen see my reply above. I was able to make it work as intended by downgrading the packages to these versions after the original versions were installed.

[tenhoeve13]

I'm still getting the error:

AttributeError: module 'jax.interpreters.pxla' has no attribute 'ShardedDeviceArray'

after trying the fix from @scploeger. Any help would be appreciated.

[tvonich]

I was above to temporarily get past this issue by downgrading the following packages:

!pip install jax==0.4.16
!pip install jaxlib==0.4.16
!pip install flax==0.7.3

This worked for me also. I only needed to downgrade jax and jaxlib

[andrewlkd]

The version of chex being pointed to is deprecated.

Running

!pip install --upgrade jaxlib !pip install --upgrade chex

should fix these issues.

Hope this helps!