Dependencies

[SamuelGabriel]

Hi There,

Thanks for open-sourcing!
The dependencies are really hard to reconstruct as it seems that Jax is changing its API quite heavily between minor releases.

How to reproduce the bug: Follow your installation manual (in my case with Python 3.9). It will install Jax 0.2 which does not work with your code :( I got it to work by changing the Jax version to 0.1.71 on the CPU but for CUDA I didn't have luck yet.

(Sorry for the tone wrote this after trying to get the dependencies right for half an hour..)

[SamuelGabriel]

Ok, I think I found a way to make it work that you could integrate into your setup.py.

I used the following versions

pip install "jax[cuda]==0.2.20" rlax==0.0.4 optax==0.0.9 chex==0.0.8 neural-tangents -f https://storage.googleapis.com/jax-releases/jax_releases.html

[iosband]

From our README:

We have tested neural_testbed on Python 3.7.

We do not guarantee that Python 3.9 will work, and I expect that it might not.
However, we should update the setup.py to pin to specific versions.

[SamuelGabriel]

I tried it out with Py 3.7 and still had the same problems on my local machine, my guess is that these problems don't appear on the Colab as there is some version of Jax installed already!? My PR (#5) is far from exhaustive in specifying version numbers for reproducibility, but it is the fix that was needed for me to be able to install it without errors on python 3.7.
There are problems with Python 3.9.7 as there is a bug in cpython (dbt-labs/dbt-core#3843 (comment)), but with Python 3.9.6, it seems to work as well now.

[mibrahimi]

Hi Samuel.

Thanks for digging into this and finding a solution. Colab does have jax preinstalled and the kernel we tested on has jax 0.2.21. I agree it is better practice to pin these packages to specific versions and we'll work on it. Please let us know if you have other issues or solutions :).

[pettepiero]

Hello,

Thanks for creating this repo.
I've been having problems getting the dependencies right. Do you maybe have a list of the packages versions that you know work together for this code?

Thank you