diff --git a/pyproject.toml b/pyproject.toml
index 67b2dea4..6763acdf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,8 +6,8 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "framesviewer>=1.0.2",
-    "jax[cuda12]>=0.5.0",
-    "jaxlib>=0.5.0",
+    "jax>=0.5.0",
+    "jax[cuda12]>=0.5.0; sys_platform == 'linux'",
     "jaxtyping>=0.2.38",
     "matplotlib>=3.10.0",
     "mediapy>=1.2.2",
@@ -16,6 +16,7 @@ dependencies = [
     "pygame>=2.6.1",
     "tensorflow>=2.18.0",
     "tf2onnx>=1.16.1",
+    "numpy<2.0",
 ]
 [build-system]
 requires = ["setuptools"]