jax version error about graphcast HOT 7 OPEN

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

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

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'