From 6894d3efef5868eaf9198b8195952ebb6752c18f Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Mon, 1 Apr 2024 08:23:59 +0000
Subject: [PATCH] Add JAX to requirements.txt

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

diff --git a/requirements-tpu.txt b/requirements-tpu.txt
index 225fbdc55d29e..6a11eebc547c7 100644
--- a/requirements-tpu.txt
+++ b/requirements-tpu.txt
@@ -4,6 +4,7 @@ sentencepiece  # Required for LLaMA tokenizer.
 numpy
 torch ~= 2.2.0
 torch_xla[tpu] ~= 2.2.0
+jax[tpu]  # Required for Pallas kernels.
 requests
 py-cpuinfo
 transformers >= 4.39.1  # Required for StarCoder2 & Llava.