Sphinx-doc for experimental.jet?

[pnkraemer]

Hi everyone,

Big fan of jet here! I came across #3730 recently, but since it might be a little out of scope for the issue, I thought I'd start a mew discussion here, instead of turning a seemingly resolved issue into a discussion. :) If ideas is the wrong category, please feel free to move the discussion.

I noticed that some experimental APIs are documented (pjit, maps, ...), while others aren't (jet, ode,...). Is there a specific rule regarding which experimental APIs are documented and which ones are not?

Because if it is a matter of someone writing the docs, I'd be very happy to contribute a sphinx-doc for jax.experimental.jet with some basic explanations and examples!

I made some notes about the API for myself recently (based on https://openreview.net/pdf?id=SkxEF3FNPH, and my own user experience), and could turn this into a pull request fairly quickly I guess. A rough draft is attached below, but I would obviously include only what you want me to include! :)

Thoughts?

[hawkinsp]

I'm sure we'd welcome a PR that includes jet in the docs. We don't mind documenting experimental things, experimental is simply a statement about the expected level of support and quality.

[jakevdp]

To answer your bolded question: no, I don't think we have any rule regarding which experimental packages are included in the HTML docs. But looking at jet.py, it's clear there hasn't yet been any effort put into module or function-level docstrings, and this is probably why it's not yet been included in the documentation build. As Peter says, a PR improving that situation would be very welcome.

[pnkraemer]

Thanks for the quick feedback @hawkinsp and @jakevdp. I'll look into it then 👍

[pnkraemer]

See #9546.