Dependencies for netket

[aashmore]

Hello everyone!

I have a question about getting netket up and running. I am trying to run it on a Paperspace instance, but running into the following problem. (I've had it running just fine on Google Colab.)

I can install netket with pip, but when trying to import I get an error. As suggested in the documentation, I've pasted the ouput of python -m netket.tools.info below (this is also the same error when trying to import netket). Thanks for looking at this - I've been tearing my hair out!

Traceback (most recent call last): File "/usr/lib/python3.9/runpy.py", line 188, in _run_module_as_main mod_name, mod_spec, code = _get_module_details(mod_name, _Error) File "/usr/lib/python3.9/runpy.py", line 111, in _get_module_details __import__(pkg_name) File "/usr/local/lib/python3.9/dist-packages/netket/__init__.py", line 24, in <module> from . import utils File "/usr/local/lib/python3.9/dist-packages/netket/utils/__init__.py", line 21, in <module> from . import _dependencies_check File "/usr/local/lib/python3.9/dist-packages/netket/utils/_dependencies_check.py", line 55, in <module> if not module_version("optax") >= (0, 1, 1): File "/usr/local/lib/python3.9/dist-packages/netket/utils/version_check.py", line 22, in module_version module = importlib.import_module(module) File "/usr/lib/python3.9/importlib/__init__.py", line 127, in import_module return _bootstrap._gcd_import(name[level:], package, level) File "/usr/local/lib/python3.9/dist-packages/optax/__init__.py", line 17, in <module> from optax import experimental File "/usr/local/lib/python3.9/dist-packages/optax/experimental/__init__.py", line 20, in <module> from optax._src.experimental.complex_valued import split_real_and_imaginary File "/usr/local/lib/python3.9/dist-packages/optax/_src/experimental/complex_valued.py", line 32, in <module> import chex File "/usr/local/lib/python3.9/dist-packages/chex/__init__.py", line 17, in <module> from chex._src.asserts import assert_axis_dimension File "/usr/local/lib/python3.9/dist-packages/chex/_src/asserts.py", line 26, in <module> from chex._src import asserts_internal as _ai File "/usr/local/lib/python3.9/dist-packages/chex/_src/asserts_internal.py", line 34, in <module> from chex._src import pytypes File "/usr/local/lib/python3.9/dist-packages/chex/_src/pytypes.py", line 25, in <module> ArraySharded = jax.interpreters.pxla.ShardedDeviceArray AttributeError: module 'jax.interpreters.pxla' has no attribute 'ShardedDeviceArray'

[PhilipVinc]

Hi @aashmore ,
I suspect that you have incompatible versions of jax and chex. That's a quite common scenario and happens because those packages don't correctly declare compatibility ranges...

The best solution is to upgrade everything to recent versions.

pip install --upgrade jax jaxlib chex optax flax numba netket

does this work?

[aashmore]

Thanks for the help with this. After much gnashing of teeth, I discovered that none of the Paperspace instances are running recent enough versions of cuda and cudnn, and they were incompatible with the required jax version.