Run sample notebooks as of 13 April 2024.

[naruya]

Thank you so much for your excellent research and for publishing the code.

As has been said in #298 and #306, some notebooks were not able to run and required downgrading of the libraries used, so I will leave a log of that for later people.

for JAX, M.D. Cookbook.ipynb

!pip install jax==0.4.23
!pip install jaxlib==0.4.23+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install -q git+https://www.github.com/google/jax-md

for Neural Network Potentials.ipynb

!pip install jax==0.4.23
!pip install jaxlib==0.4.23+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install -q git+https://www.github.com/deepmind/haiku@v0.0.11
!pip install -q git+https://www.github.com/deepmind/optax@v0.2.2
!pip install -q --upgrade git+https://www.github.com/google/jax-md

After downgrading the libraries, be careful not to install the libraries again in other cells.

Changelog for JAX
Dependencies for optax==v0.2.2
Dependencies for haiku==v0.0.11

Thanks!

[naruya]

JAX, M.D. Cookbook.ipynb

JAX.M.D.Cookbook.ipynb.-.Colab.-.Google.Chrome.2024-04-13.13-19-56.mp4

[naruya]

Neural Network Potentials.ipynb

Neural.Network.Potentials.ipynb.-.Colab.-.Google.Chrome.2024-04-13.13-57-23.mp4