pin jax/tensorflow versions

Author: sokrypton

AlphaFold2.ipynb

@@ -144,9 +144,9 @@
         "  echo \"installing colabfold...\"\n",
         "  # install dependencies\n",
         "  # We have to use \"--no-warn-conflicts\" because colab already has a lot preinstalled with requirements different to ours\n",
-        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/[email protected]\"\n",
-        "  # high risk high gain\n",
-        "  pip install -q \"jax[cuda11_cudnn805]>=0.3.8,<0.4\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
+        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/[email protected]\" \"tensorflow-cpu==2.7.4\"\n",
+        "  pip uninstall -yq jax jaxlib\n",
+        "  pip install -q \"jax[cuda]==0.3.25\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "\n",
         "  # for debugging\n",
         "  ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\n",

batch/AlphaFold2_batch.ipynb

@@ -119,9 +119,9 @@
         "if [ ! -f COLABFOLD_READY ]; then\n",
         "  # install dependencies\n",
         "  # We have to use \"--no-warn-conflicts\" because colab already has a lot preinstalled with requirements different to ours\n",
-        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/[email protected]\"\n",
-        "  # high risk high gain\n",
-        "  pip install -q \"jax[cuda11_cudnn805]>=0.3.8,<0.4\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
+        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/[email protected]\" \"tensorflow-cpu==2.7.4\"\n",
+        "  pip uninstall -yq jax jaxlib\n",
+        "  pip install -q \"jax[cuda]==0.3.25\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "  touch COLABFOLD_READY\n",
         "fi\n",
         "\n",