Pin jax and jaxlib versions

Saransh-cpp · Saransh-cpp · commit c24f3f757275 · 2021-09-27T22:58:19.000+05:30
diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -6,8 +6,8 @@ anytree >= 2.4.3
 autograd >= 1.2
 scikit-fem >= 0.2.0
 casadi >= 3.5.0
-jax>=0.1.68
-jaxlib>=0.1.47
+jax==0.2.12
+jaxlib==0.1.70
 jupyter  # For example notebooks
 pybtex
 # Note: Matplotlib is loaded for debug plots but to ensure pybamm runs
