Merge pull request #1180 from pybamm-team/fix-jax-release

pin jax version

.github/release_checklist.md

@@ -0,0 +1,4 @@
+- Increment version number in `version`
+- Increment version number in `docs/conf.py`
+- Update CHANGELOG.md with a summary of the release
+- Update (and pin) jax and jaxlib to latest version and fix any bugs that arise

requirements.txt

@@ -5,8 +5,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.1.75
+jaxlib==0.1.52
 jupyter  # For example notebooks
 # Note: Matplotlib is loaded for debug plots but to ensure pybamm runs
 # on systems without an attached display it should never be imported

setup.py

@@ -161,8 +161,8 @@ def compile_KLU():
 jax_dependencies = []
 if system() != "Windows":
     jax_dependencies = [
-        "jax>=0.1.68",
-        "jaxlib>=0.1.47",
+        "jax==0.1.75",
+        "jaxlib==0.1.52",
     ]
 
 
@@ -172,7 +172,7 @@ def compile_KLU():
 
 setup(
     name="pybamm",
-    version=load_version() + ".post1",
+    version=load_version() + ".post2",
     description="Python Battery Mathematical Modelling.",
     long_description=readme,
     long_description_content_type="text/markdown",