diff --git a/pyproject.toml b/pyproject.toml
index 53cfb6c..8ec143e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,6 +15,7 @@ dependencies = [
     "pyyaml>=6.0.1",
     "matplotlib>=3.8.2",
     "numba>=0.58.1",
+    "jax>=0.4.24",
 ]
 requires-python = ">=3.11"
 readme = "README.md"