diff --git a/pyproject.toml b/pyproject.toml
index 38550ca..c11109f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,7 +18,8 @@ dependencies = [
     "fastmcp",
     "fastapi",
     "matplotlib",
-    "Pillow"
+    "Pillow",
+    "jax"
 ]
 
 [tool.setuptools.packages.find]