Upgrade jax versions to avoid DeviceArray import issues.

PiperOrigin-RevId: 624271204
google · Apr 12, 2024 · 84520dc · 84520dc
1 parent 93b1d9f
commit 84520dc
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/requirements-jax.txt b/requirements-jax.txt
@@ -1,11 +1,11 @@
 # Optional Jax library.
-jax>=0.4.23
-jaxlib>=0.4.23
-jaxopt>=0.7
-flax>=0.6.10
-optax>=0.1.5
-chex>=0.1.7
+jax>=0.4.26
+jaxlib>=0.4.26
+jaxopt>=0.8.3
+flax>=0.8.2
+optax>=0.2.2
+chex>=0.1.86
 tfp-nightly[jax]
-equinox>=0.10.6
-jaxtyping>=0.2.20
+equinox>=0.11.3
+jaxtyping>=0.2.28
 typeguard<=2.13.3