diff --git a/pyproject.toml b/pyproject.toml
index dd7bbf2c0..5c17bf69b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -70,6 +70,7 @@ test = [
     "transformers",
     "pillow",
     "exllamav2",
+    "jax"
 ]
 serve = [
     "vllm>=0.3.0",