From 913346f624ba8d86a171dfa32a1d444cd6707a9f Mon Sep 17 00:00:00 2001
From: Agriya Khetarpal <74401230+agriyakhetarpal@users.noreply.github.com>
Date: Thu, 23 Nov 2023 00:17:44 +0530
Subject: [PATCH] #3443 Bump to latest version of `jax` and `jaxlib`

tested with `--upgrade` and `--upgrade-strategy eager` plus `--no-cache-dir`
---
 setup.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index e55036059a..87d7a06b15 100644
--- a/setup.py
+++ b/setup.py
@@ -389,8 +389,8 @@ def compile_KLU():
             # 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.18; python_version >= '3.9'",
-            "jaxlib==0.4.18; python_version >= '3.9'",
+            "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.