[DOC] add Thinking In JAX doc

[jakevdp]

Addresses #4643

Rendered page: https://jax--5426.org.readthedocs.build/en/5426/notebooks/thinking_in_jax.html

[8bitmp3]

@jakevdp Hope you didn't mind me taking a look. I've got some ML editing experience (a book, docs). Thanks for writing this Thinking 🤔 in JAX guide 👍

Update: the RST is very well-formatted btw 👏.

[jakevdp]

@levskaya pointed out in chat that the ipython-sphinx approach here has a real downside: code blocks are intermingled with IPython prompts, which makes them less convenient for copy/pasteing from the doc.

I'm going to mark this WIP for now and explore other approaches.

[8bitmp3]

@jakevdp I was wondering if there a particular reason why the team favors RST with iPython here? I understand that Sphinx works with Markdown and I noticed the following JAX guides were created as MD files: Device Memory Profiling, custom_vjp and nondiff_argnums update guide, and Profiling JAX programs.

To follow the lessons from the How to think in JAX guide, I manually converted the RST file into a Jupyter notebook paragraph by paragraph. Then, I formatted RTS-flavored code text (functions etc), so there's no ::stuff:: (but you can add links to JAX APIs manually in Markdown, I guess), ran the cells in Colab, downloaded the IPYNB file, used Jupytext to convert it into Markdown, and left it there in a raw form (with a lot of spaces that can be cleaned up by a Markdown linter): https://github.com/8bitmp3/thinking-in-jax-jakevdp-notebook/blob/main/How_to_think_in_JAX_(by_jakevdp)_notebook_version.md - if you're interested. The only thing that stands out is that cell outputs aren't clearly marked - adding Output: before e.g. jax.interpreters.xla._DeviceArray would help. (The IPYNB to Markdown process is manual but, for example, the NumPy.org docs team has integrated it into GH Actions, I think.)

Again, thanks for putting together How to think in JAX.

[jakevdp]

We've been discussing this elsewhere... the main reason I chose rst is because of the better support for doctests and cross referencing... we've had issues with our notebook-based documentation getting out-of-date. I actually wrote this content originally as a notebook, then translated to RST for that reason.

Notebooks are also very difficult to track well in version control, and require a fairly heavy toolchain to build (e.g. pandoc) so I'm leaning toward converting our existing notebook sources to something text-based, like RST or markdown. I'm not sure we're going to go the ipython-sphinx route, though, because it makes it difficult to copy code blocks (as mentioned above).

Markdown is certainly nicer than RST for composing documents, but does not support things like cross-references very well, although myst-md is a nice hybrid approach. It also doesn't do a good job of distinguishing between code blocks and output blocks.

I'm looking into myst-nb, but the doctesting story there is not particularly clear.

[dhirschfeld]

code blocks are intermingled with IPython prompts, which makes them less convenient for copy/pasteing from the doc

...unless you paste them in ipython itself which ignores them. You could also use the sphinx-copybutton extension which (IIUC) can strip the prompts.

[jakevdp]

OK, I pushed an update where I use myst-nb rather than RST for the thinking in jax content. This required removing nbsphinx in favor of myst-nb for parsing notebooks; this is pretty nice because it does so without a pandoc requirement.

I think myst-nb is a pretty good system that ticks a lot of our boxes. WDYT?

[8bitmp3]

Found the last couple of NumPys

[jakevdp]

@jakevdp does myst-nb play well with the popular JAX Quickstart and other notebooks?

Yes, you can take a look at the rendered version here: https://jax--5426.org.readthedocs.build/en/5426/notebooks/quickstart.html

[jakevdp]

Unfortunately, it looks like execution errors in myst-nb don't lead sphinx to fail, but only to a warning.

This is unfortunate because it means that failures will silently result in unexpected errors being published to our docs via cell outputs.

This issue was reported last year in executablebooks/MyST-NB#248; the suggestion there is to use the -W tag in sphinx to globally turn all warnings into errors; unfortunately that is not a good option for JAX, because documentation format issues are also reported as warnings, and we inherit many such warnings from NumPy and SciPy in our wrapped docstring.

I consider this a hard blocker: if we can't find a way to surface notebook execution errors as build failures rather than warnings, I don't think myst-nb is going to be a good solution for JAX.

[jakevdp]

Given the failed execution issue, I think it's worth pursuing other options. I'm going to revisit the idea of using RST directly.

[marcvanzee]

Interesting discussion! For our HOWTO system in Flax we are also exploring options on how to explain things with code that is tested, while still being easily usable with git. Avital suggested to use Sphinx's doctest (see google/flax#925), which is an option you mention in this thread as well.

@jakevdp: Looking at the issue you linked in myst-nb, it seems that they are planning to fix it in the coming weeks. Supposing that they fix the issue, do you think myst-nb is prefable over doctest?

[jakevdp]

I've decided to punt on the myst-nb question for now, in favor of just getting the thinking in jax content onto the website. I converted the notebook back to ipynb, reverted the old commits, and pushed the ipynb here.

[jakevdp]

Ready for review: see the final rendered doc page here: https://jax--5426.org.readthedocs.build/en/5426/notebooks/thinking_in_jax.html

[mattjj]

Btw the 'rendered page' link in the OP didn't work for me, but this does.

EDIT: the one in your most recent comment worked though.

[mattjj]

Wow this is phenomenal. Fantastic work.

Only suggestion I could come up with: maybe rephrase "JAX's layers of API" to "JAX API layering".

[jakevdp]

Only suggestion I could come up with: maybe rephrase "JAX's layers of API" to "JAX API layering".

Done

[jakevdp]

I fixed the rendered page link in the first comment (it was from an earlier iteration and was missing the notebooks subdirectory)