diff --git a/requirements-dev.txt b/requirements-dev.txt
index 0420ccbe6..dfbe51ac2 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -7,7 +7,7 @@ rich>=13.7.1
 
 # Docs site
 furo
-jax
+jax[cpu]
 matplotlib
 ml_dtypes
 myst-parser[linkify]