Upgrade to newest versions of jax + jaxlib and add Windows support for JAX Solver

.github/workflows/docker.yml

@@ -48,8 +48,7 @@ jobs:
             echo "tag=all" >> "$GITHUB_OUTPUT"
           fi
 
-      - name: Build and push Docker image to Docker Hub (no solvers)
-        if: matrix.build-args == 'No solvers'
+      - name: Build and push Docker image to Docker Hub (${{ matrix.build-args }})
         uses: docker/build-push-action@v5
         with:
           context: .
@@ -58,29 +57,5 @@ jobs:
           push: true
           platforms: linux/amd64, linux/arm64
 
-      - name: Build and push Docker image to Docker Hub (with ODES and IDAKLU solvers)
-        if: matrix.build-args == 'ODES' || matrix.build-args == 'IDAKLU'
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          file: scripts/Dockerfile
-          tags: pybamm/pybamm:${{ steps.tags.outputs.tag }}
-          push: true
-          build-args: ${{ matrix.build-args }}=true
-          platforms: linux/amd64, linux/arm64
-
-      - name: Build and push Docker image to Docker Hub (with ALL and JAX solvers)
-        if: matrix.build-args == 'ALL' || matrix.build-args == 'JAX'
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          file: scripts/Dockerfile
-          tags: pybamm/pybamm:${{ steps.tags.outputs.tag }}
-          push: true
-          build-args: ${{ matrix.build-args }}=true
-          # exclude arm64 for JAX and ALL builds for now, see
-          # https://github.com/google/jax/issues/13608
-          platforms: linux/amd64
-
       - name: List built image(s)
         run: docker images

noxfile.py

@@ -60,10 +60,9 @@ def run_coverage(session):
     """Run the coverage tests and generate an XML report."""
     set_environment_variables(PYBAMM_ENV, session=session)
     session.install("coverage", silent=False)
-    session.install("-e", ".[all]", silent=False)
+    session.install("-e", ".[all,jax]", silent=False)
     if sys.platform != "win32":
         session.install("-e", ".[odes]", silent=False)
-        session.install("-e", ".[jax]", silent=False)
     session.run("coverage", "run", "--rcfile=.coveragerc", "run-tests.py", "--nosub")
     session.run("coverage", "combine")
     session.run("coverage", "xml")
@@ -73,7 +72,7 @@ def run_coverage(session):
 def run_integration(session):
     """Run the integration tests."""
     set_environment_variables(PYBAMM_ENV, session=session)
-    session.install("-e", ".[all]", silent=False)
+    session.install("-e", ".[all,jax]", silent=False)
     if sys.platform == "linux":
         session.install("-e", ".[odes]", silent=False)
     session.run("python", "run-tests.py", "--integration")
@@ -90,10 +89,9 @@ def run_doctests(session):
 def run_unit(session):
     """Run the unit tests."""
     set_environment_variables(PYBAMM_ENV, session=session)
-    session.install("-e", ".[all]", silent=False)
+    session.install("-e", ".[all,jax]", silent=False)
     if sys.platform == "linux":
         session.install("-e", ".[odes]", silent=False)
-        session.install("-e", ".[jax]", silent=False)
     session.run("python", "run-tests.py", "--unit")
 
 
@@ -130,17 +128,16 @@ def set_dev(session):
                     external=True,
         )
     else:
-        session.run(python, "-m", "pip", "install", "-e", ".[all,dev]", external=True)
+        session.run(python, "-m", "pip", "install", "-e", ".[all,dev,jax]", external=True)
 
 
 @nox.session(name="tests")
 def run_tests(session):
     """Run the unit tests and integration tests sequentially."""
     set_environment_variables(PYBAMM_ENV, session=session)
-    session.install("-e", ".[all]", silent=False)
+    session.install("-e", ".[all,jax]", silent=False)
     if sys.platform == "linux" or sys.platform == "darwin":
         session.install("-e", ".[odes]", silent=False)
-        session.install("-e", ".[jax]", silent=False)
     session.run("python", "run-tests.py", "--all")
 
 

setup.py

@@ -265,15 +265,26 @@ def compile_KLU():
         "pandas": [
             "pandas>=1.5.0",
         ],
+        # Note: jax and jaxlib must be pinned to a specific version
+        # to avoid upstream breaking changes.
         "jax": [
-            "jax==0.4.8",
-            "jaxlib==0.4.7",
+            # 0.4.18 provides support for Jax on aarch64 containers
+            # via the PyBaMM images on Docker Hub which come with
+            # Python 3.11 installed.
+            # It also provides support for CPU-only Jax on Windows.
+            "jax==0.4.20; python_version >= '3.9'",
+            "jaxlib==0.4.20; python_version >= '3.9'",
+            # Jax 0.4.13 was the last version to support Python 3.8.
+            # Support for CPU-only Windows was added in 0.4.13, so
+            # this version supports Windows too.
+            "jax==0.4.13; python_version < '3.9'",
+            "jaxlib==0.4.13; python_version < '3.9'",
         ],
         "odes": ["scikits.odes"],
         "all": [
             "autograd>=1.6.2",
             "scikit-fem>=8.1.0",
-            "pybamm[examples,plot,cite,latexify,bpx,tqdm,pandas]"
+            "pybamm[examples,plot,cite,latexify,bpx,tqdm,pandas]",
         ],
     },
     entry_points={

tests/unit/test_expression_tree/test_operations/test_evaluate_python.py

@@ -503,7 +503,7 @@ def test_evaluator_jax(self):
         expr = pybamm.exp(a * b)
         evaluator = pybamm.EvaluatorJax(expr)
         result = evaluator(t=None, y=np.array([[2], [3]]))
-        self.assertEqual(result, np.exp(6))
+        np.testing.assert_array_almost_equal(result, np.exp(6), decimal=15)
 
         # test a constant expression
         expr = pybamm.Scalar(2) * pybamm.Scalar(3)