Tracking issue: Documentation update

[jakevdp]

Started with #18269

The plan is to create an entirely new cohesive documentation structure in docs/tutorials, to move sections of the existing docs there, and then to remove the old versions leaving HTTP redirects in their place.

Existing drafts (on the main branch) can be seen at https://jax.readthedocs.io/en/latest/tutorials/: note that this is currently unlisted in the website table of contents.

Indroductory

• Quickstart @jakevdp JAX Docs tutorials: quickstart #18613
• Key Concepts @jakevdp DOC: add key concepts doc #20347
• Thinking in JAX @8bitmp3 JAX docs tutorial: How to think in JAX 101 #18581 removed in DOC: remove copy of thinking-in-jax from new tutorial flow #20802
• Just-in-time compilation @jakevdp JAX tutorials: add intro to jit compilation #18627
• Automatic vectorization @jakevdp JAX tutorials: add automatic vectorization #18626
• Automatic differentiation @8bitmp3 JAX docs tutorial: Autodiff 101 #18676
• Debugging @8bitmp3 JAX docs tutorial: Debugging 101 #18582
• Pseudorandom numbers @jakevdp JAX tutorials: pseudorandom numbers #18631
• Working with PyTrees @mtthss JAX docs tutorial: Working with pytrees #18864
• Introduction to sharded computation @jakevdp Add sharding 101 doc #20779
• Stateful computations @jakevdp DOC: add stateful computation doc #20732
• Example: Writing a simple neural network

Intermediate

• Parallel computation
• Advanced automatic differentiation @8bitmp3
• Advanced vectorization (collectives, batched map)
• Gradient checkpointing
• Advanced debugging
• External callbacks @jakevdp JAX tutorials: add External Callbacks #18616
• Profiling and performance @dfm

Advanced

• How JAX primitives work
• Understanding Jaxprs
• Writing a custom Jaxpr interpreter
• Ahead-of-time lowering and computation
• Custom operations for GPUs with C++ and CUDA
• Pallas: a JAX kernel language

Reference

• Installation @8bitmp3 JAX docs tutorial: Installation #18580 moved to main page in DOC: update installation guide #20784

[dfm]

I'm happy to start drafting the "Profiling and performance" tutorial because it's something I'd like to learn more about unless someone else is planning on working on it!

[jakevdp]

Awesome, consider it assigned 😁