AttributeError

[tortijazz]

Currently I'm working with macOS Monterey 12.2. My Mac has an Intel Core i7 processor; recently I was trying to upgrade Netket to the current stable version (3.3.1) as the "NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems (2021)" article suggests, but does not work. When I try to import Netket, the terminal (and jupyter) displays the error:

AttributeError: module 'jax.interpreters.xla' has no attribute 'Token'

I have been searching on the internet how to fix this issue, but I haven't found anything. I tried to install other versions of jax, but the error is the same.

Anybody knows how to fix this problem?

[PhilipVinc]

Can you paste here the output of Python -m netket.tools.info?

man’s also the whole stack trace of the error you see

[tortijazz]

Yeah. I did 'pip install --upgrade pip' before upgrade Netket, but the problem persists.

[PhilipVinc]

If you give me the output of the command and the stack trace I can help you. Most likely you have an old version of mpi4jax.

[tortijazz]

This is the full output, Philip:

File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/runpy.py", line 188, in _run_module_as_main mod_name, mod_spec, code = _get_module_details(mod_name, _Error) File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/runpy.py", line 111, in _get_module_details __import__(pkg_name) File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/netket/__init__.py", line 24, in <module> from . import utils File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/netket/utils/__init__.py", line 38, in <module> from . import mpi File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/netket/utils/mpi/__init__.py", line 15, in <module> from .mpi import ( File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/netket/utils/mpi/mpi.py", line 50, in <module> import mpi4jax File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/mpi4jax/__init__.py", line 9, in <module> from ._src import ( # noqa: E402 File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/mpi4jax/_src/__init__.py", line 9, in <module> from .collective_ops.allgather import allgather # noqa: F401 File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/mpi4jax/_src/collective_ops/allgather.py", line 32, in <module> token=(type(None), xla.Token, core.Tracer), AttributeError: module 'jax.interpreters.xla' has no attribute 'Token'

[PhilipVinc]

you have a recent version of jax and an old mpi4jax.

pip install -U mpi4jax s should fix it. Otherwise paste Python -m netket.tools.info

[tortijazz]

It works!

Thank you, Philip.