From 3a7943a44a2656efb0ca2150246ab5fefb7f8fe1 Mon Sep 17 00:00:00 2001
From: Lance Wang <lancewang@google.com>
Date: Fri, 17 Jul 2026 10:19:19 -0700
Subject: [PATCH] [Tunix] Restrict JAX version in tunix/oss pyproject.toml.

PiperOrigin-RevId: 949636811
---
 .github/workflows/tpu-tests.yml | 8 ++++----
 pyproject.toml                  | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/tpu-tests.yml b/.github/workflows/tpu-tests.yml
index c8a7f8c5f..1a50f5266 100644
--- a/.github/workflows/tpu-tests.yml
+++ b/.github/workflows/tpu-tests.yml
@@ -65,7 +65,7 @@ jobs:
     - name: Install tunix dependencies
       run: |
         pip install --upgrade pip
-        pip install -e ".[prod,test]" "safetensors<0.8" --force-reinstall
+        pip install -e ".[test]" "safetensors<0.8" --force-reinstall
 
     - name: Verify TPU availability
       run: |
@@ -160,7 +160,7 @@ jobs:
           python3 -m pip install --upgrade pip setuptools wheel
 
           # Install Tunix with dev and test dependencies without overwriting the vLLM dependencies.
-          python3 -m pip install -e .[dev,test]
+          python3 -m pip install -e .[test]
           # tpu-inference/Numba needs NumPy 2.3 or less.
           python3 -m pip install numpy==2.3.5 --force-reinstall
       - name: Run tunix model tests
@@ -222,8 +222,8 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          # Reinstall Tunix with prod dependencies
-          python3 -m pip install -e .[prod] --force-reinstall
+          # Reinstall Tunix
+          python3 -m pip install -e . --force-reinstall
           python3 -m pip install --upgrade kagglehub kagglesdk
 
           # Loading tfds requires tensorflow.
diff --git a/pyproject.toml b/pyproject.toml
index e32efc0f1..9916b65ec 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,6 +26,7 @@ dependencies = [
   "hf_transfer",  # Huggingface caching in CI
   "huggingface_hub",
   "importlib_resources",
+  "jax[tpu]>=0.6.0,!=0.7.2,<0.11.0", # Jax 0.7.2 has performance regression on OSS, 0.11.0 breaks vLLM of nightly-20260406-581c4d4-f6983f0.
   "jaxtyping",
   "jinja2",  # Huggingface chat template
   "kagglehub",
@@ -59,7 +60,6 @@ docs = [
     "sphinx_contributors",
 ]
 prod = [
-  "jax[tpu]>=0.6.0,!=0.7.2", # Jax 0.7.2 has performance regression on OSS
 ]
 dev = [
   # Manully install vLLM & tpu-inferece, which depends on jax[tpu]==0.7.2