update jax minver to 0.8.1

Author: cgarciae

pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
     "numpy>=1.23.2; python_version>='3.11'",
     "numpy>=1.26.0; python_version>='3.12'",
     # keep in sync with jax-version in .github/workflows/build.yml
-    "jax>=0.7.1",
+    "jax>=0.8.1",
     "msgpack",
     "optax",
     "orbax-checkpoint",