Remove JAX library version constraint.

As of release of 0.4.37

PiperOrigin-RevId: 707011365

Author: Nush395

pyproject.toml

@@ -21,13 +21,13 @@ dependencies = [
     "typeguard>=2.13.3,<3",
     "immutabledict>=1.0.0",
     "jax>=0.4.31",
-    "jaxlib>=0.4.31,<0.4.36",
+    "jaxlib>=0.4.31",
     "jaxopt>=0.8.2",
     "flax>=0.7.5",
     "matplotlib>=3.3.0",
     "numpy>=1.24.1,<2.0.0",
     "setuptools;python_version>='3.10'",
-    "chex>=0.1.85",
+    "chex>=0.1.88",
     "equinox>=0.11.3",
     "PyYAML>=6.0.1",
     "xarray>=2024.11.0",