From 17bb3e1b9638d6eae69899a98f7d48dd54496bbb Mon Sep 17 00:00:00 2001
From: Xiaoquan Kong <u1mail2me@gmail.com>
Date: Sun, 6 Sep 2026 09:48:42 -0700
Subject: [PATCH] Gate jax[tpu] dependency behind a Linux platform marker

libtpu only ships Linux wheels, so the hard jax[tpu] dependency made
the package uninstallable on macOS and installed an unneeded TPU
runtime on other non-TPU platforms. Keep jax[tpu] on Linux and depend
on plain jax elsewhere, with the same version constraints.

Fixes #2047
---
 pyproject.toml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 76668a5303..05eac72229 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,9 @@ dependencies = [
   "importlib_resources",
   # Jax 0.7.2 has performance regression on OSS
   # Jax 0.11.1 removed jax.experimental.hijax.MutableHiType, which flax (through 0.12.8, the newest release) still subclasses
-  "jax[tpu]>=0.6.0,!=0.7.2,<0.11.1",
+  # tpu wheels are only available on linux, so we use platform marker to gate the extra
+  "jax[tpu]>=0.6.0,!=0.7.2,<0.11.1; platform_system == 'Linux'",
+  "jax>=0.6.0,!=0.7.2,<0.11.1; platform_system != 'Linux'",
   "jaxtyping",
   "jinja2",  # Huggingface chat template
   "kagglehub",