search

google-deepmind / ferminet

An implementation of the Fermionic Neural Network for ab-initio electronic structure calculations
Apache License 2.0
721 stars 120 forks source link

Jax install - issue with correct version number #43

Closed Dehavendrew closed 2 years ago

Dehavendrew commented 2 years ago

Hi,

Apologies for the basic question, but I am having issue running the test scripts provided and I believe it is an issue with the version of Jax that was installed by default within the Setup script. Below are the versions that were automatically installed:

image

And here is the error message that is causing each test to fail.

image

Currently I am running this off the main branch within a Google Colab notebook with GPU backend enabled.

What version of Jax/Jaxlib should I be using?

Thank you for the help!

jsspencer commented 2 years ago

Sorry, there's apparently an API change we hadn't picked up on. This is used deeply inside KFAC and fixing it will take some time.

In the meantime, jax 0.2.21 and jaxlib 0.17.5 work fine for me (at least on a CPU backend).