From c6ff17791e0cd7ef572b63e990bad6016ffe1074 Mon Sep 17 00:00:00 2001
From: Xingyou Song <xingyousong@google.com>
Date: Fri, 12 Apr 2024 11:38:36 -0700
Subject: [PATCH] Upgrade jax versions to avoid `DeviceArray` import issues.

PiperOrigin-RevId: 624244553
---
 requirements-jax.txt | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/requirements-jax.txt b/requirements-jax.txt
index 1992fe633..c21ae2635 100644
--- 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
\ No newline at end of file