diff --git a/.readthedocs.yml b/.readthedocs.yml
index 8a535de1f..6f807aa82 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -8,7 +8,7 @@ version: 2
 build:
   os: "ubuntu-22.04"
   tools:
-    python: "3.9"
+    python: "3.10"
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
diff --git a/docs/conf.py b/docs/conf.py
index 9eaa4ce6d..0fa4ed23e 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -67,7 +67,6 @@ extensions = [
     'sphinx.ext.napoleon',
     'sphinx.ext.viewcode',
     'matplotlib.sphinxext.plot_directive',
-    'sphinx_autodoc_typehints',
     'myst_nb',
     "sphinx_remove_toctrees",
     'sphinx_copybutton',
@@ -294,17 +293,13 @@ epub_exclude_files = ['search.html']
 
 # -- Extension configuration -------------------------------------------------
 
-# Tell sphinx-autodoc-typehints to generate stub parameter annotations including
-# types, even if the parameters aren't explicitly documented.
-always_document_param_types = True
-
-
 # Tell sphinx autodoc how to render type aliases.
+autodoc_typehints = "description"
+autodoc_typehints_description_target = "all"
 autodoc_type_aliases = {
     'ArrayLike': 'jax.typing.ArrayLike',
     'DTypeLike': 'jax.typing.DTypeLike',
 }
 
-
 # Remove auto-generated API docs from sidebars. They take too long to build.
 remove_from_toctrees = ["_autosummary/*"]
diff --git a/docs/requirements.txt b/docs/requirements.txt
index e83ea4c92..76056fcc6 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,7 +1,6 @@
 absl-py
 ipython>=8.8.0  # 8.7.0 has ipython3 lexer error
 sphinx>=7.3.2  # 7.3.0 breaks sphinx-book-theme
-sphinx-autodoc-typehints
 sphinx-book-theme>=1.0.1  # Older versions fail to pin pydata-sphinx-theme
 sphinx-copybutton>=0.5.0
 sphinx-remove-toctrees