No attribute 'KeyArray' when importing chex

[raymondchua]

Hi,
When I try to import chex, I got the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/chuaraym/cvpr/lib/python3.8/site-packages/chex/__init__.py", line 17, in <module>
    from chex._src.asserts import assert_axis_dimension
  File "/home/chuaraym/cvpr/lib/python3.8/site-packages/chex/_src/asserts.py", line 26, in <module>
    from chex._src import asserts_internal as _ai
  File "/home/chuaraym/cvpr/lib/python3.8/site-packages/chex/_src/asserts_internal.py", line 32, in <module>
    from chex._src import pytypes
  File "/home/chuaraym/cvpr/lib/python3.8/site-packages/chex/_src/pytypes.py", line 36, in <module>
    PRNGKey = jax.random.KeyArray
AttributeError: module 'jax.random' has no attribute 'KeyArray'

I am using jax version 0.2.14 and jaxlib 0.1.68. Did they cause the error?

[raymondchua]

Fixed. I manage to upgrade to latest version.