From eb0a0466a9902d5a89264fdf05503d0d9ee558e3 Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Tue, 16 Apr 2024 08:05:54 +0000
Subject: [PATCH] Add JAX requirements

---
 requirements-tpu.txt | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/requirements-tpu.txt b/requirements-tpu.txt
index 624880ffa00e4..e627bb9b9e6b5 100644
--- a/requirements-tpu.txt
+++ b/requirements-tpu.txt
@@ -1,2 +1,5 @@
 # Common dependencies
 -r requirements-common.txt
+
+jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+flax >= 0.8