diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 3bae5ec80..d31359efd 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -26,10 +26,7 @@ jobs:
       run: |
         sudo apt install -y pandoc
         python -m pip install --upgrade pip
-        pip install -r requirements.txt
-        pip install -r requirements-jax.txt
-        pip install -r requirements-test.txt
-        pip install -r docs/requirements.txt
+        pip install -r requirements.txt -r requirements-jax.txt -r docs/requirements.txt
         pip install .
     - name: Print installed dependencies
       run: |
diff --git a/docs/requirements.txt b/docs/requirements.txt
index f1b857827..a82312353 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,86 +1,10 @@
-absl-py==1.4.0
-alabaster==0.7.12
-attrs==21.2.0
-Babel==2.9.1
-backcall==0.2.0
-bleach==4.1.0
-certifi==2021.10.8
-charset-normalizer==2.0.7
-cycler==0.11.0
-decorator==5.1.0
-defusedxml==0.7.1
-docutils==0.17.1
-entrypoints==0.3
-flatbuffers==2.0
-flax==0.7.1
-idna==3.3
-imagesize==1.2.0
-importlib-resources==5.9.0
-ipykernel==5.3.4
-ipython==7.16.1
-ipython-genutils==0.2.0
-jax==0.4.13
-jaxlib==0.4.13
-jedi==0.18.0
-Jinja2==2.11.3
-jmp==0.0.4
-jsonschema==4.1.2
-jupyter-client==7.0.6
-jupyter-core==4.9.1
-jupyterlab-pygments==0.1.2
-kiwisolver==1.3.2
-latexcodec==2.0.1
-MarkupSafe==2.0.1
-matplotlib==3.3.3
-mistune==0.8.4
-nbclient==0.5.4
-nbconvert==6.2.0
-nbformat==5.1.3
-nbsphinx==0.8.9
-nest-asyncio==1.5.1
-numpy==1.24.0
-opt-einsum==3.3.0
-oset==0.1.3
-packaging==21.2
-pandas==1.3.4
-pandoc==1.0.2
-pandocfilters==1.5.0
-parso==0.8.2
-pexpect==4.8.0
-pickleshare==0.7.5
-Pillow==9.0.0
-ply==3.11
-prompt-toolkit==3.0.21
-ptyprocess==0.7.0
-pybtex==0.24.0
-pybtex-docutils==1.0.1
-Pygments==2.7.4
-pyparsing==2.4.7
-pyrsistent==0.18.0
-python-dateutil==2.8.2
-pytz==2021.3
-PyYAML==6.0
-pyzmq==23.0.0
-requests==2.26.0
-scipy==1.11.1
-seaborn==0.11.1
-six==1.16.0
-snowballstemmer==2.1.0
-sphinx==4.5.0
-sphinx-autodoc-typehints==1.11.1
+sphinx
+sphinxcontrib-bibtex
+sphinxcontrib-katex
+sphinx-autodoc-typehints
+# See https://github.com/executablebooks/sphinx-book-theme/issues/719.
 sphinx-book-theme==0.3.3
-sphinxcontrib-applehelp==1.0.2
-sphinxcontrib-bibtex==2.5.0
-sphinxcontrib-devhelp==1.0.2
-sphinxcontrib-htmlhelp==2.0.0
-sphinxcontrib-jsmath==1.0.1
-sphinxcontrib-katex==0.8.6
-sphinxcontrib-qthelp==1.0.3
-sphinxcontrib-serializinghtml==1.1.5
-tabulate==0.8.9
-testpath==0.5.0
-tornado==6.1
-traitlets==5.1.1
-urllib3==1.26.7
-wcwidth==0.2.5
-webencodings==0.5.1
+
+# Packages used for notebook execution
+nbsphinx
+ipykernel
diff --git a/requirements-test.txt b/requirements-test.txt
index 9ba6b518f..d5d7c03bd 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -1,7 +1,6 @@
 chex>=0.0.4
 cloudpickle>=1.2.2
 dill
-flax>=0.7.1
 mock>=3.0.5
 pytest-xdist
 virtualenv
diff --git a/requirements.txt b/requirements.txt
index ed929b68f..62588eb19 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,4 @@ absl-py>=0.7.1
 jmp>=0.0.2
 numpy>=1.18.0
 tabulate>=0.8.9
+flax>=0.7.1