diff --git a/pyproject.toml b/pyproject.toml
index 76668a530..05eac7222 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",