need to pin the other parts of jax as well to work on Colab, which comes with it's own version pre-installed

Author: JRJacoby

setup.cfg

@@ -41,6 +41,9 @@ install_requires =
     ipywidgets
     ipympl
     jax<0.7.0 # Required for python >3.12 which will install a version of jax that tensorflow has not updated to use properly yet
+    jax-cuda12-pjrt<0.7.0
+    jax-cuda12-plugin<0.7.0
+    jaxlib<0.7.0
 
 [options.extras_require]
 dev =