Tracking issue: inline docstrings

[jakevdp]

For many public APIs, JAX currently uses the jax._src.numpy.util.implements helper to define docstrings dynamically at runtime. This has several drawbacks:

1. The docstrings extracted from NumPy do not always precisely match the semantics of the associated JAX function. This leads to confusion for people reading these docs.
2. JAX-specific argument requirements should be documented: for example, some functions require particular parameters to be static, and we should document this.
3. The dynamic docstrings do not include example usage: we've stripped all examples from these dynamically-generated docs, because often JAX-specific logic is necessary.
4. Development environments like PyCharm are unable to show these dynamically-generated docs
5. For some submodules, the current logic necessitates runtime dependency on upstream packages (e.g. jax.scipy imports scipy), and this adds to the import-time startup cost.
6. Wrapped docstrings refer to symbols at import time, so if downstream libraries remove a symbol, JAX cannot be imported! This is currently causing problems for jax versions earlier than 0.4.27, because recent scipy removed symbols from its namespace.
7. The implements decorator, despite being correctly annotated, can cause some type checkers to not recognize the input/output annotations of the wrapped function (this has been observed with pytype)

The implements decorator was originally added as a way to quickly bootstrap development of the jax.numpy API, but that is no longer necessary. For all these reasons, we'd like to migrate to static, JAX-specific inline docstrings.

A challenge here is that we cannot simply copy and modify the existing dynamic docstrings: they come from numpy, and for licensing reasons, we cannot replicate NumPy docstring verbiage within the JAX repository. For this reason, each docstring must be rewritten from scratch, without referencing the text of the NumPy version of the documentation.

[apghml]

Looking forward to it!

they come from numpy, and for licensing reasons, we cannot replicate NumPy docstring verbiage within the JAX repository

IANAL, but isn't Numpy under a 3-clause BSD license, which is Apache-compatible? The Apache Software Foundation itself seems to allow 3-clause BSD code within Apache Software Foundation projects licensed under the Apache license: https://www.apache.org/legal/resolved.html

[jakevdp]

IANAL either, which is why I do not specualate on such things 😁. But for reasons I won't get into, rewriting from scratch is the way we need to do it here.