From 876ab692cf1289f6ba9fc893a9da426b8164905d Mon Sep 17 00:00:00 2001
From: Andy Tan <andytan@google.com>
Date: Wed, 2 Sep 2026 00:30:18 -0700
Subject: [PATCH] Relax JAX version pins and add jax[cuda12] to and-cuda extra

PiperOrigin-RevId: 974951353
---
 CHANGELOG.md   |  3 +++
 pyproject.toml | 11 ++++++-----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ce108c97d..14cb34827 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,6 +23,9 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+*   Relax `jax` and `jaxlib` dependency pins to `>= 0.7.2, < 1.0.0` and add
+    `jax[cuda12]` to `[and-cuda]` optional dependencies to support Python 3.13
+    and prevent PJRT accelerator plugin version mismatches.
 *   Update `BayesianPPPCheck` calculation to use the posterior predictive
     distribution with `sigma` rather than the posterior expected outcome.
 *   **Breaking change**: Toggle default backend from TensorFlow to JAX.
diff --git a/pyproject.toml b/pyproject.toml
index 54de26567..5433e92e6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,8 +29,8 @@ dependencies = [
   "babel",
   "bidict",
   "immutabledict",
-  "jax == 0.7.2",
-  "jaxlib == 0.7.2",
+  "jax >= 0.7.2, < 1.0.0",
+  "jaxlib >= 0.7.2, < 1.0.0",
   "joblib",
   "matplotlib < 3.11.0",
   "natsort >= 7.1.1, < 8",
@@ -71,6 +71,7 @@ colab = [
 # Installed through `pip install -e .[and-cuda]`
 and-cuda = [
   "tensorflow[and-cuda] >= 2.21.0, < 2.22",
+  "jax[cuda12] >= 0.7.2, < 1.0.0",
 ]
 # MLflow deps
 # Installed through `pip install -e .[mlflow]`
@@ -78,9 +79,9 @@ mlflow = ["mlflow"]
 # JAX backend dependencies.
 # Installed through `pip install -e .[jax]`
 jax = [
-    "jax==0.7.2",
-    "jaxlib==0.7.2",
-    "tfp-nightly[substrates-jax]==0.26.0.dev20260130",
+  "jax >= 0.7.2, < 1.0.0",
+  "jaxlib >= 0.7.2, < 1.0.0",
+  "tfp-nightly[substrates-jax] == 0.26.0.dev20260130",
 ]
 # MMM proto schema.
 schema = [