From 909fee6a2d4b0cdb70bb3c85d5cb9a317996c596 Mon Sep 17 00:00:00 2001
From: Matthew Johnson <mattjj@google.com>
Date: Mon, 30 Mar 2020 20:22:04 -0700
Subject: [PATCH] try adding sphinx-autodoc-typehints

---
 docs/conf.py          | 1 +
 docs/requirements.txt | 1 +
 2 files changed, 2 insertions(+)

diff --git a/docs/conf.py b/docs/conf.py
index 90f8a9239..9ef10a9dd 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -60,6 +60,7 @@ extensions = [
     'sphinx.ext.napoleon',
     'sphinx.ext.viewcode',
     'nbsphinx',
+    'sphinx_autodoc_typehints',
 ]
 
 intersphinx_mapping = {
diff --git a/docs/requirements.txt b/docs/requirements.txt
index ac8fd2a93..c7f76dce0 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -3,6 +3,7 @@ sphinx_rtd_theme
 jaxlib
 ipykernel
 nbsphinx
+sphinx-autodoc-typehints
 # The next packages are for notebooks
 matplotlib
 sklearn