bug: multipledispatch error

Hello,

I’ve updagraded jax to 0.3.5, and now I am getting the following error: ModuleNotFoundError: No module named 'multipledispatch'. To reproduce: python -c "import gpjax".

After pip install multipledispatch, I get the another error:

ImportError: Encountered error: `Tried to dispatch on non-type: typing.Union[numpy.ndarray, jax._src.numpy.ndarray.ndarray]
In signature: <typing.Union[numpy.ndarray, jax._src.numpy.ndarray.ndarray], Distribution>
In function: log_density`

Perhaps a specific version of multipledispatch is expected but it’s in the requirements.txt?

pip freeze ->

absl-py==1.0.0
aiohttp==3.8.1
aiosignal==1.2.0
antlr4-python3-runtime==4.8
asttokens==2.0.5
astunparse==1.6.3
async-timeout==4.0.2
attrs==21.4.0
autograd==1.3
backcall==0.2.0
black==22.1.0
blessings==1.7
cachetools==5.0.0
certifi==2021.10.8
charset-normalizer==2.0.12
chex==0.1.3
click==8.0.4
cloudpickle==2.0.0
colorlog==6.6.0
contextlib2==21.6.0
cycler==0.11.0
debugpy==1.5.1
decorator==5.1.1
dill==0.3.4
distrax==0.1.2
dm-haiku==0.0.6
dm-tree==0.1.6
docker-pycreds==0.4.0
e3nn-jax==0.6.2
einops==0.4.1
entrypoints==0.4
etils==0.6.0
executing==0.8.3
flatbuffers==2.0
fonttools==4.30.0
frozenlist==1.3.0
fsspec==2022.2.0
future==0.18.2
gast==0.5.3
-e git+ssh://[e-mail]/oxcsml/geomstats.git@78111276c9b2f98bdc24826f0c2d8729a4bbcca0#egg=geomstats&subdirectory=../../geomstats
gitdb==4.0.9
GitPython==3.1.27
google-auth==2.6.0
google-auth-oauthlib==0.4.6
google-pasta==0.2.0
googleapis-common-protos==1.55.0
GPJax==0.4.1
gpustat==0.6.0
grpcio==1.44.0
h5py==3.6.0
hydra-colorlog==1.1.0
hydra-core==1.1.1
hydra-joblib-launcher==1.1.5
hydra-submitit-launcher @ git+https://github.com/emilemathieu/hydra.git@7c41ae8f080d577efccf902aa9c955d7c9a2d864#subdirectory=plugins/hydra_submitit_launcher
idna==3.3
importlib-metadata==4.11.2
importlib-resources==5.8.0
install==1.3.5
ipykernel==6.9.1
ipython==8.1.1
jax==0.3.5
jaxlib==0.3.5+cuda11.cudnn805
jaxtyping==0.0.2
jedi==0.18.1
Jinja2==3.1.2
jmp==0.0.2
joblib==1.1.0
jupyter-client==7.1.2
jupyter-core==4.9.2
keras==2.8.0
Keras-Preprocessing==1.1.2
kiwisolver==1.3.2
libclang==13.0.0
Markdown==3.3.6
MarkupSafe==2.1.1
matplotlib==3.5.1
matplotlib-inline==0.1.3
ml-collections==0.1.0
mpmath==1.2.1
multidict==6.0.2
multipledispatch==0.6.0
mypy-extensions==0.4.3
nest-asyncio==1.5.4
numpy==1.22.3
nvidia-ml-py3==7.352.0
oauthlib==3.2.0
omegaconf==2.1.1
opt-einsum==3.3.0
optax==0.1.1
packaging==21.3
pandas==1.4.1
parso==0.8.3
pathspec==0.9.0
pathtools==0.1.2
pexpect==4.8.0
pickleshare==0.7.5
Pillow==9.0.1
platformdirs==2.5.1
promise==2.3
prompt-toolkit==3.0.28
protobuf==3.19.0
psutil==5.9.0
ptyprocess==0.7.0
pure-eval==0.2.2
pyasn1==0.4.8
pyasn1-modules==0.2.8
Pygments==2.11.2
pyparsing==3.0.7
python-dateutil==2.8.2
pytz==2021.3
PyYAML==6.0
pyzmq==22.3.0
requests==2.27.1
requests-oauthlib==1.3.1
rsa==4.8
scikit-learn==1.0.2
scipy==1.8.0
-e git+ssh://[e-mail]/oxcsml/score-sde.git@384fa52c4ad3a6125dea4af61178b5ad9db65dfd#egg=score_sde
seaborn==0.11.2
sentry-sdk==1.5.7
setGPU==0.0.7
setproctitle==1.2.2
shortuuid==1.0.8
six==1.16.0
smmap==5.0.0
stack-data==0.2.0
submitit @ git+https://github.com/emilemathieu/submitit.git@aeccb447a1e1a33ec60749c70594eef3d888b16d
sympy==1.10.1
tabulate==0.8.9
tensorboard==2.8.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow==2.8.1
tensorflow-datasets==4.5.2
tensorflow-estimator==2.8.0
tensorflow-io-gcs-filesystem==0.24.0
tensorflow-metadata==1.7.0
tensorflow-probability==0.15.0
termcolor==1.1.0
tf-estimator-nightly==2.8.0.dev2021122109
threadpoolctl==3.1.0
tomli==2.0.1
toolz==0.11.2
tornado==6.1
tqdm==4.63.0
traitlets==5.1.1
typeguard==2.13.3
typing_extensions==4.3.0
urllib3==1.26.8
wandb==0.12.11
wcwidth==0.2.5
Werkzeug==2.0.3
wrapt==1.14.0
yarl==1.7.2
yaspin==2.1.0
zipp==3.7.0

[thomaspinder]

Good to know - thanks for the information @Daniel-Dodd and @holmrenser. When I get some time, I’ll type this up and include in our docs as you’ll no doubt not be the last people to experience M1/Jax/GPJax installationissues.

Best of luck @holmrenser! I’d love to hear how you get on - any feedback is greatly appreciated either through an issue, discussion thread or PR 😃

[holmrenser]

@Daniel-Dodd this is a fix for me as well!

I get some warnings, and first time running takes 10-20seconds. After that it’s a bit faster.

(gpjax) rensholmer@Renss-MBP-2 ~ % python3 -c 'import gpjax;print(gpjax.__version__)'
/Users/rensholmer/miniconda3/envs/gpjax/lib/python3.10/site-packages/jax/_src/lib/__init__.py:33: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
  warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
0.4.9

Off to experimenting with gpjax!