`jax.random` causes build fail

[geo-179]

Python version: 3.10

Python package versions:

absl-py==2.1.0
astunparse==1.6.3
cachetools==5.3.3
certifi==2024.2.2
charset-normalizer==3.3.2
chex==0.1.85
contourpy==1.2.0
cycler==0.12.1
distlib==0.3.8
etils==1.7.0
filelock==3.13.1
flatbuffers==23.5.26
flax==0.8.1
fonttools==4.49.0
fsspec==2024.2.0
gast==0.4.0
gin==0.1.6
gin-config==0.1.1
google-auth==2.28.1
google-auth-oauthlib==1.0.0
google-pasta==0.2.0
grpcio==1.62.0
h5py==3.10.0
huggingface-hub==0.21.3
idna==3.6
importlib_resources==6.1.2
jax==0.4.25
jaxlib==0.4.25
keras==2.13.1
kiwisolver==1.4.5
libclang==16.0.6
Markdown==3.5.2
markdown-it-py==3.0.0
MarkupSafe==2.1.5
matplotlib==3.8.3
mdurl==0.1.2
ml-dtypes==0.3.2
msgpack==1.0.8
nest-asyncio==1.6.0
numpy==1.24.3
oauthlib==3.2.2
opt-einsum==3.3.0
optax==0.1.9
orbax-checkpoint==0.5.3
packaging==23.2
pillow==10.2.0
platformdirs==4.2.0
protobuf==4.25.3
pyasn1==0.5.1
pyasn1-modules==0.3.0
Pygments==2.17.2
pyparsing==3.1.1
python-dateutil==2.9.0.post0
PyYAML==6.0.1
regex==2023.12.25
requests==2.31.0
requests-oauthlib==1.3.1
rich==13.7.1
rsa==4.9
safetensors==0.4.2
scipy==1.12.0
six==1.16.0
tensorboard==2.13.0
tensorboard-data-server==0.7.2
tensorflow==2.13.0
tensorflow-estimator==2.13.0
tensorflow-macos==2.13.0
tensorstore==0.1.54
termcolor==2.4.0
tokenizers==0.15.2
toolz==0.12.1
tqdm==4.66.2
transformers==4.38.2
typing_extensions==4.5.0
urllib3==2.2.1
virtualenv==20.25.1
Werkzeug==3.0.1
wrapt==1.16.0
zipp==3.17.0

To build the package, I simply used ./run.sh. If I try running ./run_tests.sh, I get AttributeError: module 'jax.random' has no attribute 'KeyArray'

@0xWOLAND