From 2cf4d354f6f0b723628c6d6e225417febfebd933 Mon Sep 17 00:00:00 2001
From: Yu Li <wstcliyu@google.com>
Date: Thu, 24 Sep 2026 16:33:19 +0000
Subject: [PATCH] Add JAX 0.11.2 to test matrix and bump orbax-checkpoint to
 0.12.3

---
 .github/workflows/test.yml | 4 +++-
 pyproject.toml             | 2 +-
 requirements.txt           | 2 +-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e4538cd..7e2be80 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -27,12 +27,14 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ['3.11', '3.12', '3.13', '3.14']
-        jax-version: ['0.8.3', '0.9.0', '0.9.1', '0.9.2', '0.10.0', '0.10.1', '0.10.2', '0.11.0', '0.11.1', 'nightly']
+        jax-version: ['0.8.3', '0.9.0', '0.9.1', '0.9.2', '0.10.0', '0.10.1', '0.10.2', '0.11.0', '0.11.1', '0.11.2', 'nightly']
         exclude:
           - python-version: '3.11'
             jax-version: '0.11.0'
           - python-version: '3.11'
             jax-version: '0.11.1'
+          - python-version: '3.11'
+            jax-version: '0.11.2'
 
     runs-on: ubuntu-latest
     timeout-minutes: 30
diff --git a/pyproject.toml b/pyproject.toml
index a0c95da..ca7676a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ dependencies = [
     "absl-py",
     "fastapi",
     "jax>=0.8.3",
-    "orbax-checkpoint",
+    "orbax-checkpoint>=0.12.3",
     "uvicorn",
     "requests",
     "packaging",
diff --git a/requirements.txt b/requirements.txt
index c03f8df..98dd1d9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@ absl-py
 fastapi
 httpx
 jax>=0.8.0
-orbax-checkpoint
+orbax-checkpoint>=0.12.3
 uvicorn
 requests
 packaging