Repository: google/jax

Issue:
docstring processing no longer fully working
e.g. 
```python
>>> print('versionadded' in jnp.nansum.__doc__)                                                                                    
True
```
In the past this was removed here: https://github.com/google/jax/blob/97a5719fcb40af7231b5f803f965063538282f8e/jax/_src/numpy/util.py#L127-L132

But something has changed and that is no longer happening.


Discussion:
I think the `jit` decorators removed `__code__` (https://github.com/google/jax/blob/04f322e065a2cdbb988bf73f26b6ae4968bf1ea5/jax/_src/numpy/util.py#L124)

Resolving patch:
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -121,14 +121,15 @@ def wrap(op):
       try:
         parsed = _parse_numpydoc(docstr)
 
-        if update_doc and hasattr(op, '__code__') and 'Parameters' in parsed.sections:
+        if update_doc and 'Parameters' in parsed.sections:
+          code = getattr(getattr(op, "__wrapped__", op), "__code__", None)
           # Remove unrecognized parameter descriptions.
           parameters = _parse_parameters(parsed.sections['Parameters'])
           parsed.sections['Parameters'] = (
             "Parameters\n"
             "----------\n" +
             "\n".join(_versionadded.split(desc)[0].rstrip() for p, desc in parameters.items()
-                      if p in op.__code__.co_varnames and p not in skip_params)
+                      if (code is None or p in code.co_varnames) and p not in skip_params)
           )
 
         docstr = parsed.summary.strip() + "\n" if parsed.summary else ""


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -44,7 +44,7 @@
 from jax import tree_util
 from jax.test_util import check_grads
 from jax._src.util import prod, safe_zip
-from jax._src.numpy.util import _parse_numpydoc, ParsedDoc
+from jax._src.numpy.util import _parse_numpydoc, ParsedDoc, _wraps
 from jax._src.numpy.lax_numpy import _promote_dtypes, _promote_dtypes_inexact
 
 from jax.config import config
@@ -6207,6 +6207,46 @@ def test_lax_numpy_docstrings(self):
       if obj.__doc__ and "*Original docstring below.*" not in obj.__doc__:
         raise Exception(f"jnp.{name} does not have a wrapped docstring.")
 
+  @parameterized.named_parameters(
+    {"testcase_name": "_jit" if jit else "", "jit": jit} for jit in [True, False])
+  def test_wrapped_function_parameters(self, jit):
+    def orig(x):
+      """Example Docstring
+
+      Parameters
+      ----------
+      x : array_like
+        Input Data
+
+        .. versionadded:: 1.8.0
+      out : array_like, optional
+        Output to overwrite
+      other_arg : Any
+        not used
+
+      Returns
+      -------
+      x : input
+      """
+      return x
+
+    def wrapped(x, out=None):
+      return x
+
+    if jit:
+      wrapped = jax.jit(wrapped)
+
+    wrapped = _wraps(orig, skip_params=['out'])(wrapped)
+    doc = wrapped.__doc__
+
+    self.assertStartsWith(doc, "Example Docstring")
+    self.assertIn("Original docstring below", doc)
+    self.assertIn("Parameters", doc)
+    self.assertIn("Returns", doc)
+    self.assertNotIn('out', doc)
+    self.assertNotIn('other_arg', doc)
+    self.assertNotIn('versionadded', doc)
+
 
   def test_parse_numpydoc(self):
     # Unit test ensuring that _parse_numpydoc correctly parses docstrings for all