From c473fdee633a74a1d4e44d305ae8f0cd47690d8d Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Thu, 1 Apr 2021 13:27:11 -0700
Subject: [PATCH] docs: tweaks to make docs build on Python 3.8

---
 .readthedocs.yml            |  2 +-
 docs/jax.profiler.rst       | 10 ++++++++++
 jax/_src/numpy/lax_numpy.py |  6 ++++--
 jax/_src/numpy/util.py      |  5 +++++
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 1e5dc20b6c84..92e8df0b170c 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -16,6 +16,6 @@ formats:
 
 # Optionally set the version of Python and requirements required to build your docs
 python:
-  version: 3.7
+  version: 3.8
   install:
     - requirements: docs/requirements.txt
diff --git a/docs/jax.profiler.rst b/docs/jax.profiler.rst
index 7e99fc38879e..c1f61af1cd67 100644
--- a/docs/jax.profiler.rst
+++ b/docs/jax.profiler.rst
@@ -34,3 +34,13 @@ profiling features.
 
   device_memory_profile
   save_device_memory_profile
+
+Deprecated functions
+--------------------
+
+.. autosummary::
+  :toctree: _autosummary
+
+  trace_function
+  TraceContext
+  StepTraceContext
\ No newline at end of file
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 7d8b4a434448..87a1cc8cef1e 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1928,10 +1928,12 @@ def _isposneginf(infinity, x, out):
     return full_like(x, False, dtype=bool_)
 
 isposinf = _wraps(np.isposinf, skip_params=['out'])(
-  lambda x, out=None: _isposneginf(inf, x, out))
+  lambda x, out=None: _isposneginf(inf, x, out)
+)
 
 isneginf = _wraps(np.isneginf, skip_params=['out'])(
-  lambda x, out=None: _isposneginf(-inf, x, out))
+  lambda x, out=None: _isposneginf(-inf, x, out)
+)
 
 @_wraps(np.isnan)
 def isnan(x):
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
index 38547ed2cea9..d5fcdd93c243 100644
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -21,6 +21,7 @@
 _section_break = re.compile(r"\n(?=[^\n]{3,15}\n-{3,15})", re.MULTILINE)
 _numpy_signature_re = re.compile(r'^([\w., ]+=)?\s*[\w\.]+\([\w\W]*?\)$', re.MULTILINE)
 _versionadded = re.compile(r'^\s+\.\.\s+versionadded::', re.MULTILINE)
+_docreference = re.compile(r':doc:`(.*?)\s*<.*?>`')
 
 class ParsedDoc(NamedTuple):
   """
@@ -48,6 +49,10 @@ def _parse_numpydoc(docstr: Optional[str]) -> ParsedDoc:
   if docstr is None or not docstr.strip():
     return ParsedDoc(docstr)
 
+  # Remove any :doc: directives in the docstring to avoid sphinx errors
+  docstr = _docreference.sub(
+    lambda match: f"{match.groups()[0]}", docstr)
+
   signature, body = "", docstr
   match = _numpy_signature_re.match(body)
   if match: