lower jax requirement

Author: n-gao

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
 ]
 dependencies = [
-    "jax>=0.4.30; python_version < '3.13'",
+    "jax>=0.4.10; python_version < '3.13'",
     "jax>=0.4.34; python_version >= '3.13'",
     "jaxtyping>=0.2.36",
     "numpy>=1.25; python_version < '3.13'",
@@ -53,7 +53,7 @@ build-backend = "hatchling.build"
 [dependency-groups]
 dev = [
     "flax>=0.10.2",
-    "jax[cuda12]>=0.4.38",
+    "jax[cuda12]>=0.4.10",
     "parameterized>=0.9.0",
     "pre-commit-uv>=4.1.4",
     "pytest>=8.3.4",