Remove dependency on sphinx_autodoc_typehints

.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:

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/*"]

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