docs: tweaks to make docs build on Python 3.8

.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

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

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

jax/_src/numpy/util.py

@@ -21,6 +21,7 @@ _parameter_break = re.compile("\n(?=[A-Za-z_])")
 _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: