From 2af2f19790c93c1d025aca9c278517be6ae44c9d Mon Sep 17 00:00:00 2001
From: Tom Hennigan <tomhennigan@google.com>
Date: Fri, 16 Feb 2024 06:38:44 -0800
Subject: [PATCH] Use tf-nightly for Haiku testing.

PiperOrigin-RevId: 607673638
---
 requirements-jax.txt  | 4 ++--
 requirements-test.txt | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/requirements-jax.txt b/requirements-jax.txt
index a305fd19e..683585e48 100644
--- a/requirements-jax.txt
+++ b/requirements-jax.txt
@@ -1,2 +1,2 @@
-jax>=0.4.16
-jaxlib>=0.4.16
+jax>=0.4.24
+jaxlib>=0.4.24
diff --git a/requirements-test.txt b/requirements-test.txt
index d5d7c03bd..b868ebe32 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -5,5 +5,5 @@ mock>=3.0.5
 pytest-xdist
 virtualenv
 dm-tree>=0.1.1
-tensorflow>=2.9.1
+tf-nightly>=2.16.0
 optax>=0.0.1