JAX tutorials draft

Note

The tutorials below are a work in progress; for the time being, please refer to the older tutorial content, including Getting Started with JAX, JAX 101, and the now-deleted JAX 101 tutorials.

JAX 101

Mostly finalized at JAX 101!

• Key concepts
• Just-in-time compilation
• Automatic vectorization
• Automatic differentiation
• Introduction to debugging
• Pseudorandom numbers
• Pytrees
• Introduction to parallel programming
• Stateful computations
• Example: Writing a simple neural network

JAX 201

• Parallel computation
• Advanced debugging
• Profiling and performance

JAX 301

• Advanced compilation