WIP: Refactor "How to think in JAX" and surrounding pages

[melissawm]

Starting with the idea of reorganizing some of the information in the docs, this PR is a first attempt at moving content to make each page's scope more defined.

This is very much a work in progress, and more discussion is needed before an actual review.

A few thoughts:

• Although the initial proposal was to move all the "How to think in JAX" contents into the "Key concepts" page, I think they should be kept separate, but with clearer audiences. "How to think in JAX" is a more practical, hands-on and introductory (like a tutorial or how-to), while "Key concepts" is more of an explanation, addressing the why.
• I see a progression from readers from Quickstart -> How to think in JAX -> Key concepts so I proposed a reorg for these pages, but kept all three at the top level as it seems appropriate to keep them more visible for new readers.
• Some of the contents in these pages reference "previous" or "later" sections and should be reworded to use more cross links and less positional wording.
• I'd like to verify that the Key concepts ToC is representative of the Key concepts we want to mention.
• There is some repetition in contents (e.g. "Jax vs. NumPy", Pytrees) across pages. This is appropriate, as long as each page reaches a different audience.
• Why is jax.core not documented? There are a few mentions of jax.core.Tracer that can't be linked since no documentation is generated for that module or class. I've added this in this PR, but since there are no docstrings I'm not sure this is the best option.

Addresses #24632

cc @emilyfertig

[emilyfertig]

In theory I like the idea that you laid out in the first couple bullets, of having a flow from a practical/hands-on HTTIJ to a more explanatory "Key Concepts" section. My concern is that the material we currently have doesn't support it well.

I agree with moving the section on jax.Array to HTTIJ. The "Pseudorandom numbers" section also reads like a practical how-to, so I'd propose moving that too. "Transformations" is already covered in the Quickstart (in slightly greater detail), so I'm not sure that section adds much (although it is a key concept). The remaining sections are "Tracing," "Jaxprs", and "Pytrees." Jaxprs feels a little too esoteric to get this top billing (one could happily use JAX for a very long time without worrying about them, IMO), and Pytrees could also fit well in a more practically-focused section. That leaves Tracing on its own (one of the other bullets in the original proposal was to rework Tracing into its own tutorial).

I think where I lean, given the material that's there, is to stick with the original proposal and consolidate these. I do like the idea of practical/explanatory intros, so if we could rearrange/add some other material in that direction, I can see that working well too. Curious what others think. cc @jakevdp

[melissawm]

Thanks, @emilyfertig - I am looking more deeply into this now. One extra thing to note is that the "How to think in JAX" document is a notebook that can be opened in Colab, so it would work nicely as an interactive introduction/quickstart. Alternatively, we could transform the quickstart into an interactive document and continue with the original plan of moving the contents into the Key Concepts page.

[jakevdp]

Some history here: "How to think in JAX" came first, but when we reorganized the getting started docs last year, we wrote "key concepts" and planned on deleting "how to think in JAX" because it's duplicative. But someone (I forgot who) argued that it's still useful and has a memorable name, so we opted to keep it, but leave it separate from the main introductory flow.

We can always revisit past decisions of course, but I thought that might be helpful context.

[melissawm]

Here is another take: in this commit, I merged the "How to think in JAX" document with the Quickstart. Here's my rationale:

• The How to Think in JAX page is very practical, in that is points out specific techniques for using JAX effectively. It feels relevant for folks coming from NumPy or a general array programming mindset.
• I wanted to make sure the Quickstart was executable, as I think being able to open it in Colab and playing with it is very useful for new users.

Again, this is just experimentation and I can instead work with a less radical approach, keeping both the How to think in JAX and Quickstart separate. I'd still argue for both to be executable, though.

[emilyfertig]

I really like what you did with this! I think merging some of the HTTIJ into Quickstart works well, and I agree that the short jax.Array section fits better there (and it's a nice segue from the "JAX vs. NumPy" section that already hints at jax.Arrays). The consolidation you've done makes a lot of sense to me.

Have you run pre-commit? I think that should catch the errors like the ones the presubmit picked up.

Somewhat related to my other comment, this isn't a strongly-held view, but I think that a good eventual end-state might be removing Key Concepts (probably not in this PR), or else fleshing it out into something that feels more coherent.

[melissawm]

Thanks @emilyfertig - I think I got your comments 🙏🏻

[melissawm]

Done @emilyfertig !