diff --git a/docs/requirements.txt b/docs/requirements.txt
index 890e508..613ebfc 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -7,6 +7,7 @@ jupyter-sphinx>=0.3.2
 matplotlib
 myst-nb
 numpy
+optax
 scikit-learn
 sphinx>=5
 sphinx-autodoc-typehints