From afef1e8eb109321eefc9ac0cd4b1cc637feae95e Mon Sep 17 00:00:00 2001
From: Jonathan Calderon Chavez <chavezcalderon@google.com>
Date: Fri, 4 Apr 2025 17:10:38 +0000
Subject: [PATCH] bump jaxs

---
 tpu/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tpu/requirements.txt b/tpu/requirements.txt
index ec6fb273..72e61d1d 100644
--- a/tpu/requirements.txt
+++ b/tpu/requirements.txt
@@ -12,7 +12,7 @@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TOR
 torchaudio==${TORCHAUDIO_VERSION}
 torchvision==${TORCHVISION_VERSION}
 # Jax packages
-jax[tpu]>=0.4.34
+jax[tpu]>=0.5.3
 --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 distrax
 flax